remove tapeagents dep, add task-level obs postprocess

Author: ollmer

src/agentlab/actions.py

@@ -1,9 +1,11 @@
 import json
 import logging
+from typing import Literal
 
 from bgym import AbstractActionSet
-from tapeagents.tool_calling import FunctionCall, ToolCallAction, ToolSpec
+from pydantic import BaseModel, Field
 
+from agentlab.backends.browser.base import FunctionCall, ToolCallAction, ToolSpec
 from agentlab.llm.llm_utils import parse_html_tags_raise
 
 logger = logging.getLogger(__name__)

src/agentlab/agents/generic_agent/generic_agent.py

@@ -10,13 +10,10 @@
 
 from copy import deepcopy
 from dataclasses import asdict, dataclass
-from functools import partial
 from warnings import warn
 
-import bgym
 from bgym import Benchmark
 from browsergym.experiments.agent import Agent, AgentInfo
-from tapeagents.tool_calling import ToolSpec
 
 from agentlab.actions import ToolsActionSet
 from agentlab.agents import dynamic_prompting as dp
@@ -67,7 +64,7 @@ def prepare(self):
     def close(self):
         return self.chat_model_args.close_server()
 
-    def make_agent(self, actions: list[ToolSpec] | None = None):
+    def make_agent(self, actions: list | None = None):
         return GenericAgent(
             chat_model_args=self.chat_model_args,
             flags=self.flags,
@@ -83,7 +80,7 @@ def __init__(
         chat_model_args: BaseModelArgs,
         flags: GenericPromptFlags,
         max_retry: int = 4,
-        actions: list[ToolSpec] | None = None,
+        actions: list | None = None,
     ):
 
         self.chat_llm = chat_model_args.make_model()

src/agentlab/backends/browser/env.py

@@ -1,31 +1,15 @@
-import json
 import logging
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal
-
-from tapeagents.core import Action, Observation, StopStep
-from tapeagents.tool_calling import ToolCallAction, ToolSpec
 
 from agentlab.actions import ToolsActionSet
-from agentlab.backends.browser.base import BrowserBackend
+from agentlab.backends.browser.base import BrowserBackend, ToolCallAction, ToolSpec
 from agentlab.benchmarks.abstract_env import AbstractEnv, AbstractEnvArgs
 from agentlab.benchmarks.web_task import AbstractWebTask
 
 logger = logging.getLogger(__name__)
 
-
-class GoalObservation(Observation):
-    kind: Literal["goal_observation"] = "goal_observation"
-    goal: str
-
-
-class PageObservation(Observation):
-    kind: Literal["page_observation"] = "page_observation"
-    content: str
-
-
 class BrowserEnv(AbstractEnv):
     def __init__(
         self, task_name: str, task: AbstractWebTask, backend: BrowserBackend, seed: int = 0
@@ -50,22 +34,23 @@ def reset(self, seed: int):
         page_content = self.backend.page_snapshot()
         screenshot = self.backend.page_screenshot()
         logger.info(f"Initial obs: {page_content}\n{screenshot}")
-        return {
+        obs = {
             "goal_object": [{"type": "text", "text": self.goal}],
             "pruned_html": page_content,
             "axtree_txt": page_content,
             "screenshot": screenshot,
             "last_action_error": "",
             "focused_element_bid": "none",
-        }, {}
+        }
+        return self.task.obs_postprocess(obs), {}
 
-    def step(self, action: ToolCallAction | str) -> tuple[Observation, float, bool, bool, dict]:
+    def step(self, action: ToolCallAction | str) -> tuple[dict, float, bool, bool, dict]:
         if isinstance(action, str):
             action = ToolsActionSet.parse_action(action)
         logger.info(f"BrowserEnv.step() called with action {action}")
 
         action_exec_start = time.time()
-        finished = isinstance(action, StopStep)
+        finished = action.function.name == "final_step"
         if finished:
             observation = {
                 "goal_object": [{"type": "text", "text": self.goal}],
@@ -76,6 +61,7 @@ def step(self, action: ToolCallAction | str) -> tuple[Observation, float, bool,
             }
         else:
             observation = self._step(action)
+        observation = self.task.obs_postprocess(observation)
         action_exec_stop = time.time()
         self._turns += 1
         logger.info(f"Obs:\n{observation['pruned_html']}")
@@ -95,8 +81,7 @@ def step(self, action: ToolCallAction | str) -> tuple[Observation, float, bool,
             "action_exec_stop": action_exec_stop,
             "action_exec_timeout": 0.0,
         } | other
-        obs_view = observation.short_view() if isinstance(observation, Observation) else observation
-        logger.info(f"Action result in observation: {obs_view}")
+        logger.info(f"Action result in observation: {observation}")
         return observation, reward, finished, truncated, env_info
 
     def _step(self, action: ToolCallAction) -> dict:
@@ -108,7 +93,7 @@ def _step(self, action: ToolCallAction) -> dict:
             "focused_element_bid": "none",
         }
 
-    def validate_task(self, action: Action, observation: PageObservation) -> tuple[float, dict]:
+    def validate_task(self, action: ToolCallAction, observation: dict) -> tuple[float, dict]:
         validate_js = self.task.get_step_validate_js()
         validate_result = self.backend.run_js(validate_js)
         reward, other = self.task.parse_validation_result(validate_result)

src/agentlab/backends/browser/mcp_playwright.py

@@ -2,11 +2,9 @@
 import logging
 from io import BytesIO
 
-from mcp.types import ImageContent, TextContent
 from PIL import Image
-from tapeagents.tool_calling import ToolCallAction
 
-from agentlab.backends.browser.mcp import MCPBrowserBackend
+from agentlab.backends.browser.mcp import MCPBrowserBackend, ToolCallAction
 
 logger = logging.getLogger(__name__)
 

src/agentlab/benchmarks/miniwob/task.py

@@ -3,6 +3,7 @@
 from typing import Any, ClassVar
 
 from browsergym.miniwob import ALL_MINIWOB_TASKS
+from PIL import Image
 
 from agentlab.benchmarks.web_task import AbstractWebTask
 
@@ -130,6 +131,12 @@ def parse_validation_result(self, validation_result: str) -> tuple[float, dict]:
             "done": done,
         }
 
+    def obs_postprocess(self, obs: dict) -> dict:
+        screenshot: Image.Image | None = obs.get("screenshot", None)
+        if screenshot is not None:
+            obs["screenshot"] = screenshot.crop((0, 0, 332, 214)) # crop to 332x214 because this is the viewport size for MiniWob
+        return obs
+
 
 def get_miniwob_tasks(
     base_url: str | None = None, remove_human_display: bool = True, episode_max_time: int = 1000000

src/agentlab/benchmarks/web_task.py

@@ -1,7 +1,8 @@
 from typing import ClassVar
 
 from pydantic import BaseModel
-from tapeagents.tool_calling import ToolSpec
+
+from agentlab.backends.browser.base import ToolSpec
 
 
 class AbstractWebTask(BaseModel):
@@ -29,3 +30,6 @@ def get_step_validate_js(self) -> str:
 
     def parse_validation_result(self, validate_result: str) -> tuple[float, dict]:
         raise NotImplementedError
+
+    def obs_postprocess(self, obs: dict) -> dict:
+        return obs