From cc1bf28c7e57d244023ed389c5aa892a9f9e8b46 Mon Sep 17 00:00:00 2001 From: Oleh Shliazhko Date: Wed, 15 Oct 2025 18:27:27 +0000 Subject: [PATCH 1/3] mock chat class inside the browsergym browser --- examples/rl_webagent/environment.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/examples/rl_webagent/environment.py b/examples/rl_webagent/environment.py index d1d356a48..da4e522d9 100644 --- a/examples/rl_webagent/environment.py +++ b/examples/rl_webagent/environment.py @@ -21,6 +21,23 @@ WebTask, ) +from browsergym.core.chat import Chat +# Mock the Chat class to avoid huge slowdown caused by it +def mock_chat_init(self, *args, **kwargs): + self.messages = [] + +def mock_wait_for_user_message(self, *args, **kwargs): + pass + +def mock_add_message(self, role: str, msg: str): + self.messages.append({"role": role, "timestamp": time.time(), "message": msg}) + +Chat.__init__ = mock_chat_init +Chat.wait_for_user_message = mock_wait_for_user_message +Chat.add_message = mock_add_message + + + logger = logging.getLogger(__name__) From 875b56c916f8451b9ff998f389876c35ccca24b3 Mon Sep 17 00:00:00 2001 From: Oleh Shliazhko Date: Mon, 27 Oct 2025 17:20:48 +0000 Subject: [PATCH 2/3] mock chat close, log mock use and record action time --- examples/rl_webagent/environment.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/examples/rl_webagent/environment.py b/examples/rl_webagent/environment.py index da4e522d9..e477f47ed 100644 --- a/examples/rl_webagent/environment.py +++ b/examples/rl_webagent/environment.py @@ -22,23 +22,32 @@ ) 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 = logging.getLogger(__name__) +logger.info("Mocked Chat class initialized") class WebEnvironment(Environment): @@ -187,9 +196,11 @@ def react(self, tape: WebTape) -> WebTape: # 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], ...]: From 9ec04d1ce7ce5ba1b4b2a74c33c23ca60ce8f59a Mon Sep 17 00:00:00 2001 From: Oleh Shliazhko Date: Mon, 3 Nov 2025 16:35:19 +0000 Subject: [PATCH 3/3] claude tool use --- tapeagents/llms/claude.py | 44 ++++++++++++++++++++++++++++++++------- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/tapeagents/llms/claude.py b/tapeagents/llms/claude.py index 6f32ac7e0..16fb0d179 100644 --- a/tapeagents/llms/claude.py +++ b/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)