Speed up web env task start

examples/rl_webagent/environment.py

@@ -1,28 +1,54 @@
 import logging
 import os
 import time
 from typing import Any, Literal
 
 from browsergym.core.task import AbstractBrowserTask
 from browsergym.miniwob import ALL_MINIWOB_TASKS
 from browsergym.miniwob.base import AbstractMiniwobTask
 
 from tapeagents.core import Action, FinalObservation, LLMOutputParsingFailureAction, Observation
 from tapeagents.environment import Environment
 from tapeagents.steps import ActionExecutionFailure
 from tapeagents.tools.browser import Browser
 from tapeagents.utils import FatalError
 
 from .steps import (
     FinalAnswerAction,
     ReflectionThought,
     WebTape,
     WebTapeMetadata,
     WebTask,
 )
 
+from browsergym.core.chat import Chat
+
 logger = logging.getLogger(__name__)
 
+# Mock the Chat class to avoid huge slowdown caused by it
+def mock_chat_init(self, *args, **kwargs):
+    logger.info("Mocked Chat.__init__")
+    self.messages = []
+
+def mock_wait_for_user_message(self, *args, **kwargs):
+    logger.info("Mocked Chat.wait_for_user_message")
+    pass
+
+def mock_add_message(self, role: str, msg: str):
+    logger.info("Mocked Chat.add_message")
+    self.messages.append({"role": role, "timestamp": time.time(), "message": msg})
+
+def mock_close(self, *args, **kwargs):
+    logger.info("Mocked Chat.close")
+    pass
+
+Chat.__init__ = mock_chat_init
+Chat.wait_for_user_message = mock_wait_for_user_message
+Chat.add_message = mock_add_message
+Chat.close = mock_close
+
+logger.info("Mocked Chat class initialized")
+
 
 class WebEnvironment(Environment):
     """
@@ -170,9 +196,11 @@
         # TODO: MAYBE make sure to update parent_id, author_name, etc... in the new tape.metadata just like in agent.run()
 
     def step(self, action: Action) -> Observation:
+        t = time.perf_counter()
         obs = self.browser.run(action)
         if obs.metadata.other.get("env_finished", False):
             obs = FinalObservation(metadata=obs.metadata)
+        obs.metadata.other["action_execution_time"] = time.perf_counter() - t
         return obs
 
     def actions(self) -> tuple[type[Action], ...]:

tapeagents/llms/claude.py

@@ -3,6 +3,7 @@
 from typing import Generator
 
 import anthropic
+from litellm import ChatCompletionMessageToolCall
 from omegaconf import DictConfig, OmegaConf
 
 from tapeagents.core import Prompt
@@ -12,6 +13,13 @@
 
 logger = logging.getLogger(__name__)
 
+def as_anthropic_tool(tool_spec_dict: dict) -> dict:
+    return {
+        "name": tool_spec_dict["function"]["name"],
+        "description": tool_spec_dict["function"]["description"],
+        "input_schema": tool_spec_dict["function"]["parameters"],
+    }
+
 
 class Claude(CachedLLM):
     max_tokens: int = 4096
@@ -33,13 +41,25 @@ def _generate(
         messages = self.update_image_messages_format(messages)
         while True:
             try:
-                response: anthropic.types.Message = anthropic.Anthropic().messages.create(
-                    model=self.model_name,
-                    max_tokens=self.max_tokens,
-                    system=system_message,
-                    messages=messages,
-                    **kwargs,
-                )
+                if prompt.tools is not None and len(prompt.tools) > 0:
+                    tools = [as_anthropic_tool(tool) for tool in prompt.tools]
+                    logger.info(f"Tools: {tools}")
+                    response: anthropic.types.Message = anthropic.Anthropic().messages.create(
+                        model=self.model_name,
+                        max_tokens=self.max_tokens,
+                        system=system_message,
+                        messages=messages,
+                        tools=tools,
+                        **kwargs,
+                    )
+                else:
+                    response: anthropic.types.Message = anthropic.Anthropic().messages.create(
+                        model=self.model_name,
+                        max_tokens=self.max_tokens,
+                        system=system_message,
+                        messages=messages,
+                        **kwargs,
+                    )
                 break
             except anthropic.RateLimitError as e:
                 retry_count += 1
@@ -66,7 +86,15 @@ def _generate(
                 output = LLMOutput(content=content_block.text)
                 yield LLMEvent(output=output)
             elif content_block.type == "tool_use":
-                output = LLMOutput(tool_calls=[content_block])
+                logger.info(f"Tool use: {content_block}")
+                tool_call = ChatCompletionMessageToolCall(
+                    id=content_block.id,
+                    function=dict(
+                        name=content_block.name,
+                        arguments=content_block.input,
+                    ),
+                )
+                output = LLMOutput(tool_calls=[tool_call])
                 yield LLMEvent(output=output)
             elif content_block.type == "thinking":
                 output = LLMOutput(content=content_block.text)