simpler tool call object

Author: ollmer

src/agentlab/actions.py

@@ -1,10 +1,11 @@
 import json
 import logging
-from typing import Any, Callable, Literal
+from typing import Callable, Literal
+from uuid import uuid4
 
 from bgym import AbstractActionSet
 from langchain_core.utils.function_calling import convert_to_openai_tool
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from agentlab.llm.llm_utils import parse_html_tags_raise
 
@@ -26,22 +27,11 @@ class FunctionSpec(BaseModel):
     parameters: dict
 
 
-class FunctionCall(BaseModel):
-    """
-    A class representing a function call.
-
-    Attributes:
-        name (str): The name of the function being called.
-        arguments (Any): The arguments to be passed to the function.
-    """
 
+class ToolCall(BaseModel):
+    id: str = Field(default_factory=lambda: uuid4().hex)
     name: str
-    arguments: Any
-
-
-class ToolCallAction(BaseModel):
-    id: str = ""
-    function: FunctionCall
+    arguments: dict = Field(default_factory=dict)
 
     def llm_view(self, **kwargs) -> str:
         return self.model_dump_json(indent=2)
@@ -111,7 +101,7 @@ def example_action(self, abstract: bool) -> str:
 }"""
 
     @classmethod
-    def parse_action(cls, llm_output: str) -> ToolCallAction:
+    def parse_action(cls, llm_output: str) -> ToolCall:
         logger.info(f"Parsing action: {llm_output}")
         if "<action>" in llm_output:
             content_dict, valid, retry_message = parse_html_tags_raise(llm_output, keys=["action"])
@@ -124,7 +114,7 @@ def parse_action(cls, llm_output: str) -> ToolCallAction:
             action_dict = json.loads(action_str)
         except json.JSONDecodeError:
             raise ValueError(f"Failed to parse action: {action_str}")
-        return ToolCallAction(function=FunctionCall(name=action_dict["name"], arguments=action_dict["arguments"]))
+        return ToolCall(name=action_dict["name"], arguments=action_dict["arguments"])
 
     def to_python_code(self, action) -> str:
         return action

src/agentlab/agents/react_toolcall_agent.py

@@ -10,7 +10,7 @@
 from PIL import Image
 from termcolor import colored
 
-from agentlab.actions import FunctionCall, ToolCallAction, ToolsActionSet, ToolSpec
+from agentlab.actions import ToolCall, ToolsActionSet, ToolSpec
 from agentlab.agents.agent_args import AgentArgs
 from agentlab.llm.chat_api import BaseModelArgs
 from agentlab.llm.llm_utils import image_to_png_base64_url
@@ -112,15 +112,13 @@ def obs_to_messages(self, obs: dict) -> list[dict]:
                 )
         return messages
 
-    def get_action(self, obs: dict) -> tuple[ToolCallAction, dict]:
+    def get_action(self, obs: dict) -> tuple[ToolCall, dict]:
         actions_count = len(
             [msg for msg in self.history if isinstance(msg, Message) and msg.tool_calls]
         )
         if actions_count >= self.config.max_actions:
             logger.warning("Max actions reached, stopping agent.")
-            stop_action = ToolCallAction(
-                id="stop", function=FunctionCall(name="final_step", arguments={})
-            )
+            stop_action = ToolCall(name="final_step")
             return stop_action, {}
         self.history += self.obs_to_messages(self.obs_preprocessor(obs))
         tools = [tool.model_dump() for tool in self.action_set.actions]
@@ -155,15 +153,14 @@ def thoughts_from_message(self, message) -> str:
             thoughts.append(message.content)
         return "\n\n".join(thoughts)
 
-    def action_from_message(self, message) -> ToolCallAction:
+    def action_from_message(self, message) -> ToolCall:
         if message.tool_calls:
             if len(message.tool_calls) > 1:
                 logger.warning("Multiple tool calls found in LLM response, using the first one.")
             tool_call = message.tool_calls[0]
+            name = tool_call.function.name
             args = json.loads(tool_call.function.arguments)
-            action = ToolCallAction(
-                id=tool_call.id, function=FunctionCall(name=tool_call.function.name, arguments=args)
-            )
+            action = ToolCall(id=tool_call.id, name=name, arguments=args)
             self.last_tool_call_id = action.id
             logger.info(f"Parsed tool call action: {action}")
         else:

src/agentlab/backends/browser/__init__.py

@@ -1,4 +1,3 @@
-from agentlab.actions import FunctionCall, ToolCallAction, ToolSpec
 from agentlab.backends.browser.base import BrowserBackend
 from agentlab.backends.browser.env import BrowserEnv, BrowserEnvArgs
 from agentlab.backends.browser.mcp import MCPBrowserBackend, MCPClient
@@ -7,9 +6,6 @@
 
 __all__ = [
     "BrowserBackend",
-    "FunctionCall",
-    "ToolCallAction",
-    "ToolSpec",
     "BrowserEnv",
     "BrowserEnvArgs",
     "MCPBrowserBackend",

src/agentlab/backends/browser/base.py

@@ -4,7 +4,7 @@
 from PIL import Image
 from pydantic import BaseModel
 
-from agentlab.actions import ToolCallAction, ToolSpec
+from agentlab.actions import ToolCall, ToolSpec
 
 logger = logging.getLogger(__name__)
 
@@ -35,7 +35,7 @@ def page_axtree(self) -> str:
         pass
 
     @abstractmethod
-    def step(self, action: ToolCallAction) -> str:
+    def step(self, action: ToolCall) -> dict:
         pass
 
     @abstractmethod

src/agentlab/backends/browser/env.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from agentlab.actions import ToolCallAction, ToolsActionSet, ToolSpec
+from agentlab.actions import ToolCall, ToolsActionSet, ToolSpec
 from agentlab.backends.browser.base import BrowserBackend
 from agentlab.benchmarks.abstract_env import AbstractEnv, AbstractEnvArgs
 from agentlab.benchmarks.web_task import AbstractWebTask
@@ -52,13 +52,13 @@ def reset(self, seed: int):
         obs = self.task.obs_postprocess(obs)
         return obs, {}
 
-    def step(self, action: ToolCallAction | str) -> tuple[dict, float, bool, bool, dict]:
+    def step(self, action: ToolCall | 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 = action.function.name == "final_step"
+        finished = action.name == "final_step"
         if finished:
             observation = {
                 "goal_object": [{"type": "text", "text": self.goal}],
@@ -91,7 +91,7 @@ def step(self, action: ToolCallAction | str) -> tuple[dict, float, bool, bool, d
         logger.info(f"Action result in observation: {observation}")
         return observation, reward, finished, truncated, env_info
 
-    def _step(self, action: ToolCallAction) -> dict:
+    def _step(self, action: ToolCall) -> dict:
         obs_dict = self.backend.step(action)
         if "goal_object" not in obs_dict:
             obs_dict["goal_object"] = [{"type": "text", "text": self.goal}]
@@ -101,7 +101,7 @@ def _step(self, action: ToolCallAction) -> dict:
             obs_dict["focused_element_bid"] = "none"
         return obs_dict
 
-    def validate_task(self, action: ToolCallAction, observation: dict) -> tuple[float, dict]:
+    def validate_task(self, action: ToolCall, 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.py

@@ -10,7 +10,7 @@
 from mcp import Tool as MCPTool
 from mcp.types import CallToolResult, ImageContent, TextContent
 
-from agentlab.actions import FunctionSpec, ToolCallAction, ToolSpec
+from agentlab.actions import FunctionSpec, ToolCall, ToolSpec
 from agentlab.backends.browser.base import BrowserBackend
 
 logger = logging.getLogger(__name__)
@@ -150,13 +150,12 @@ def initialize(self) -> None:
         self._mcp = MCPClient(config_path=self.config_path)
         self._mcp.initialize()
 
-    def step(self, action: ToolCallAction) -> dict:
-        contents = self.call_tool(action.function.name, action.function.arguments)
+    def step(self, action: ToolCall) -> dict:
+        contents = self.call_tool(action.name, action.arguments)
         text = "\n".join([c.text for c in contents if c.type == "text"])
         images = [c for c in contents if c.type == "image"]
         return {
-            "pruned_html": text,
-            "axtree_txt": text,
+            "text": text,
             "screenshot": images[-1] if images else None,
         }
 

src/agentlab/backends/browser/mcp_playwright.py

@@ -4,7 +4,7 @@
 
 from PIL import Image
 
-from agentlab.actions import ToolCallAction
+from agentlab.actions import ToolCall
 from agentlab.backends.browser.mcp import MCPBrowserBackend
 
 logger = logging.getLogger(__name__)
@@ -27,8 +27,8 @@ def run_js(self, js: str):
             raise e
         return result_str
 
-    def step(self, action: ToolCallAction) -> dict:
-        contents = self.call_tool(action.function.name, action.function.arguments)
+    def step(self, action: ToolCall) -> dict:
+        contents = self.call_tool(action.name, action.arguments)
         logger.info(f"Step result has {len(contents)} contents")
         tool_result = "\n".join(
             [c.text for c in contents if c.type == "text" and "# Ran Playwright code" not in c.text]

src/agentlab/backends/browser/playwright.py

@@ -6,7 +6,7 @@
 from PIL import Image
 from playwright.async_api import Browser, Page, async_playwright
 
-from agentlab.actions import ToolCallAction, ToolSpec
+from agentlab.actions import ToolCall, ToolSpec
 from agentlab.backends.browser.base import BrowserBackend
 
 logger = logging.getLogger(__name__)
@@ -106,13 +106,13 @@ def page_axtree(self):
         flat_axtree = flatten_axtree(axtree)
         return flat_axtree
 
-    def step(self, action: ToolCallAction):
-        fn = self._actions[action.function.name]
+    def step(self, action: ToolCall):
+        fn = self._actions[action.name]
         try:
-            action_result = self._loop.run_until_complete(fn(**action.function.arguments))
+            action_result = self._loop.run_until_complete(fn(**action.arguments))
         except Exception as e:
-            logger.error(f"Error executing action {action.function.name}: {e}")
-            action_result = f"Error executing action {action.function.name}: {e}"
+            action_result = f"Error executing action {action.name}: {e}"
+            logger.error(action_result)
         html = self.page_html()
         screenshot = self.page_screenshot()
         axtree = self.page_axtree()