openenv/browsergym_env

1"""Harness-oriented BrowserGym session adapters."""2 3from __future__ import annotations4 5import ast6import json7from typing import Any8 9from openenv.core.env_server.mcp_types import Tool10from openenv.core.harness import (11    ResourceSessionFactory,12    StepEnvSessionAdapter,13    ToolResult,14    VerifyResult,15)16from openenv.core.llm_client import ToolCall17 18from .models import BrowserGymAction19 20_BROWSERGYM_TOOLS = [21    Tool(22        name="click",23        description="Click an element by BrowserGym bid.",24        input_schema={25            "type": "object",26            "properties": {27                "bid": {"type": "string"},28            },29            "required": ["bid"],30        },31    ),32    Tool(33        name="fill",34        description="Fill an input field by bid.",35        input_schema={36            "type": "object",37            "properties": {38                "bid": {"type": "string"},39                "text": {"type": "string"},40            },41            "required": ["bid", "text"],42        },43    ),44    Tool(45        name="send_keys",46        description="Send keyboard input to the page.",47        input_schema={48            "type": "object",49            "properties": {50                "text": {"type": "string"},51            },52            "required": ["text"],53        },54    ),55    Tool(56        name="scroll",57        description="Scroll the page up or down.",58        input_schema={59            "type": "object",60            "properties": {61                "direction": {"type": "string", "enum": ["up", "down"]},62            },63            "required": ["direction"],64        },65    ),66    Tool(67        name="noop",68        description="Take no action on the current page.",69        input_schema={"type": "object", "properties": {}},70    ),71]72 73 74def _quote(value: str) -> str:75    return json.dumps(value, ensure_ascii=False)76 77 78def build_browsergym_action_str(tool_name: str, arguments: dict[str, Any]) -> str:79    """Convert a BrowserGym tool call into the action string the env expects."""80 81    if tool_name == "click":82        return f"click({_quote(str(arguments['bid']))})"83    if tool_name == "fill":84        return (85            f"fill({_quote(str(arguments['bid']))}, {_quote(str(arguments['text']))})"86        )87    if tool_name == "send_keys":88        return f"send_keys({_quote(str(arguments['text']))})"89    if tool_name == "scroll":90        return f"scroll({_quote(str(arguments['direction']))})"91    if tool_name == "noop":92        return "noop()"93 94    raise KeyError(f"Unsupported BrowserGym tool: {tool_name}")95 96 97def _format_browsergym_prompt(observation: Any, task: Any) -> str:98    goal = getattr(observation, "goal", "") or (task or "")99    page_text = getattr(observation, "axtree_txt", "") or getattr(100        observation, "text", ""101    )102    error = getattr(observation, "error", "")103 104    parts = []105    if goal:106        parts.append(f"Goal: {goal}")107    if error:108        parts.append(f"Previous action error: {error}")109    if page_text:110        parts.append(f"Page structure:\n{page_text}")111    parts.append("Choose the next browser action.")112    return "\n\n".join(parts)113 114 115def _build_browsergym_tool_result(116    task: Any,117):118    def builder(119        tool_name: str,120        arguments: dict[str, Any],121        result: Any,122        state: Any,123    ) -> ToolResult:124        observation = result.observation125        data = {126            "tool_name": tool_name,127            "arguments": dict(arguments),128            "goal": getattr(observation, "goal", "") or (task or ""),129            "observation_text": getattr(observation, "axtree_txt", "")130            or getattr(observation, "text", ""),131            "url": getattr(observation, "url", ""),132            "error": getattr(observation, "error", ""),133            "last_action_error": getattr(observation, "last_action_error", False),134            "reward": result.reward,135            "done": result.done,136        }137        metadata = {138            "reward": result.reward,139            "state": state.model_dump() if hasattr(state, "model_dump") else state,140        }141        return ToolResult(data=data, done=bool(result.done), metadata=metadata)142 143    return builder144 145 146def _build_browsergym_verify(147    transcript: list[dict[str, Any]],148    final_state: Any | None,149    last_result: Any | None,150    state: Any,151) -> VerifyResult:152    reward = None if last_result is None else last_result.reward153    done = False if last_result is None else bool(last_result.done)154    metrics = {155        "step_count": getattr(state, "step_count", 0),156        "cum_reward": getattr(state, "cum_reward", reward or 0.0),157        "benchmark": getattr(state, "benchmark", ""),158        "task_name": getattr(state, "task_name", ""),159    }160    artifacts = {161        "final_state": state.model_dump() if hasattr(state, "model_dump") else state,162        "final_rollout": final_state,163        "transcript_length": len(transcript),164    }165    return VerifyResult(166        env_reward=reward,167        done=done,168        metrics=metrics,169        artifacts=artifacts,170    )171 172 173class BrowserGymSessionFactory(ResourceSessionFactory):174    """Create BrowserGym-backed resource sessions from client factories."""175 176    def __init__(self, client_factory, *, default_task: str | None = None):177        self._client_factory = client_factory178        self._default_task = default_task179 180    def create(181        self,182        task: Any,183        seed: int | None = None,184        episode_id: str | None = None,185    ) -> StepEnvSessionAdapter:186        session_task = task if task is not None else self._default_task187        client = self._client_factory()188 189        reset_kwargs = {}190        if session_task is not None:191            reset_kwargs["task_name"] = session_task192 193        return StepEnvSessionAdapter(194            client=client,195            task=session_task,196            seed=seed,197            episode_id=episode_id,198            tool_specs=list(_BROWSERGYM_TOOLS),199            action_builder=lambda name, arguments: BrowserGymAction(200                action_str=build_browsergym_action_str(name, arguments)201            ),202            initial_messages_builder=lambda result, current_task: [203                {204                    "role": "user",205                    "content": _format_browsergym_prompt(206                        result.observation,207                        current_task,208                    ),209                }210            ],211            tool_result_builder=_build_browsergym_tool_result(session_task),212            verify_builder=_build_browsergym_verify,213            reset_kwargs=reset_kwargs,214        )215 216 217def _parse_action_call(action_text: str) -> tuple[str, list[Any]]:218    try:219        expression = ast.parse(action_text.strip(), mode="eval").body220    except SyntaxError as exc:221        raise ValueError(f"Unsupported BrowserGym action: {action_text}") from exc222 223    if not isinstance(expression, ast.Call) or not isinstance(224        expression.func, ast.Name225    ):226        raise ValueError(f"Unsupported BrowserGym action: {action_text}")227    if expression.keywords:228        raise ValueError("BrowserGym action arguments must be positional")229 230    args: list[Any] = []231    for arg in expression.args:232        try:233            args.append(ast.literal_eval(arg))234        except (SyntaxError, ValueError) as exc:235            raise ValueError("BrowserGym action arguments must be literals") from exc236 237    return expression.func.id, args238 239 240def _expect_str(value: Any, argument_name: str) -> str:241    if not isinstance(value, str):242        raise ValueError(f"BrowserGym {argument_name} argument must be a string")243    return value244 245 246def build_browsergym_action_tool_call(action_text: str) -> ToolCall:247    """Parse a text BrowserGym action into a structured tool call."""248 249    name, args = _parse_action_call(action_text)250    if name == "click" and len(args) == 1:251        return ToolCall(252            id="browsergym-click",253            name="click",254            args={"bid": _expect_str(args[0], "bid")},255        )256    if name == "fill" and len(args) == 2:257        return ToolCall(258            id="browsergym-fill",259            name="fill",260            args={261                "bid": _expect_str(args[0], "bid"),262                "text": _expect_str(args[1], "text"),263            },264        )265    if name == "send_keys" and len(args) == 1:266        return ToolCall(267            id="browsergym-send_keys",268            name="send_keys",269            args={"text": _expect_str(args[0], "text")},270        )271    if name == "scroll" and len(args) == 1:272        direction = _expect_str(args[0], "direction")273        if direction not in {"up", "down"}:274            raise ValueError("BrowserGym scroll direction must be 'up' or 'down'")275        return ToolCall(276            id="browsergym-scroll",277            name="scroll",278            args={"direction": direction},279        )280    if name == "noop" and not args:281        return ToolCall(id="browsergym-noop", name="noop", args={})282 283    raise ValueError(f"Unsupported BrowserGym action: {action_text}")284 285 286__all__ = [287    "BrowserGymSessionFactory",288    "build_browsergym_action_str",289    "build_browsergym_action_tool_call",290]291