diff --git a/src/agentlab/agents/agent_args.py b/src/agentlab/agents/agent_args.py index 78036c319..9928b2431 100644 --- a/src/agentlab/agents/agent_args.py +++ b/src/agentlab/agents/agent_args.py @@ -1,3 +1,5 @@ +from dataclasses import dataclass + import bgym from bgym import AbstractAgentArgs, Benchmark @@ -45,3 +47,20 @@ def set_reproducibility_mode(self): raise NotImplementedError( f"set_reproducibility_mode is not implemented for agent_args {self.__class__.__name__}" ) + + +@dataclass +class ChatModelAgentArgs(AgentArgs): + """Base class for agents whose LLM is configured via ``chat_model_args``. + + Provides the ``prepare``/``close`` implementations shared by every agent + that holds a ``chat_model_args`` factory, starting and stopping the + underlying LLM server. Subclasses are still responsible for declaring the + ``chat_model_args`` field and implementing ``make_agent``. + """ + + def prepare(self): + return self.chat_model_args.prepare_server() + + def close(self): + return self.chat_model_args.close_server() diff --git a/src/agentlab/agents/agent_utils.py b/src/agentlab/agents/agent_utils.py index 29219d2d1..7a170e8ab 100644 --- a/src/agentlab/agents/agent_utils.py +++ b/src/agentlab/agents/agent_utils.py @@ -1,7 +1,52 @@ +from dataclasses import asdict + +from browsergym.experiments.agent import AgentInfo from PIL import Image, ImageDraw from playwright.sync_api import Page +def busted_retry_ans_dict(max_retry: int) -> dict: + """Fallback answer dict used when parsing the LLM response keeps failing. + + Args: + max_retry: The maximum number of retries the agent allowed. + + Returns: + An answer dict with a null action and the retry bookkeeping marking the + step as a busted retry. + """ + return dict(action=None, n_retry=max_retry + 1, busted_retry=1) + + +def make_agent_info(chat_llm, ans_dict: dict, chat_messages, chat_model_args) -> AgentInfo: + """Build the ``AgentInfo`` returned by chat-based agents from ``get_action``. + + Collects the LLM call stats, augments them with the retry bookkeeping found + in ``ans_dict`` and packages everything (thought, messages, stats, model + args) into an ``AgentInfo``. + + Args: + chat_llm: The chat model, used to retrieve per-call stats. + ans_dict: The parsed answer dict (must contain ``n_retry`` and + ``busted_retry``; ``think`` is optional). + chat_messages: The messages exchanged with the model for this step. + chat_model_args: The model args dataclass, serialized into ``extra_info``. + + Returns: + The populated ``AgentInfo``. + """ + stats = chat_llm.get_stats() + stats["n_retry"] = ans_dict["n_retry"] + stats["busted_retry"] = ans_dict["busted_retry"] + + return AgentInfo( + think=ans_dict.get("think", None), + chat_messages=chat_messages, + stats=stats, + extra_info={"chat_model_args": asdict(chat_model_args)}, + ) + + def draw_mouse_pointer(image: Image.Image, x: int, y: int) -> Image.Image: """ Draws a semi-transparent mouse pointer at (x, y) on the image. diff --git a/src/agentlab/agents/generic_agent/generic_agent.py b/src/agentlab/agents/generic_agent/generic_agent.py index d1f48f76e..c9754fbad 100644 --- a/src/agentlab/agents/generic_agent/generic_agent.py +++ b/src/agentlab/agents/generic_agent/generic_agent.py @@ -9,16 +9,17 @@ """ from copy import deepcopy -from dataclasses import asdict, dataclass +from dataclasses import dataclass from functools import partial from warnings import warn import bgym from bgym import Benchmark -from browsergym.experiments.agent import Agent, AgentInfo +from browsergym.experiments.agent import Agent from agentlab.agents import dynamic_prompting as dp -from agentlab.agents.agent_args import AgentArgs +from agentlab.agents.agent_args import ChatModelAgentArgs +from agentlab.agents.agent_utils import busted_retry_ans_dict, make_agent_info from agentlab.llm.chat_api import BaseModelArgs from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry from agentlab.llm.tracking import cost_tracker_decorator @@ -27,7 +28,7 @@ @dataclass -class GenericAgentArgs(AgentArgs): +class GenericAgentArgs(ChatModelAgentArgs): chat_model_args: BaseModelArgs = None flags: GenericPromptFlags = None max_retry: int = 4 @@ -59,12 +60,6 @@ def set_benchmark(self, benchmark: Benchmark, demo_mode): def set_reproducibility_mode(self): self.chat_model_args.temperature = 0 - def prepare(self): - return self.chat_model_args.prepare_server() - - def close(self): - return self.chat_model_args.close_server() - def make_agent(self): return GenericAgent( chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry @@ -134,16 +129,8 @@ def get_action(self, obs): ans_dict["busted_retry"] = 0 # inferring the number of retries, TODO: make this less hacky ans_dict["n_retry"] = (len(chat_messages) - 3) / 2 - except ParseError as e: - ans_dict = dict( - action=None, - n_retry=self.max_retry + 1, - busted_retry=1, - ) - - stats = self.chat_llm.get_stats() - stats["n_retry"] = ans_dict["n_retry"] - stats["busted_retry"] = ans_dict["busted_retry"] + except ParseError: + ans_dict = busted_retry_ans_dict(self.max_retry) self.plan = ans_dict.get("plan", self.plan) self.plan_step = ans_dict.get("step", self.plan_step) @@ -151,12 +138,7 @@ def get_action(self, obs): self.memories.append(ans_dict.get("memory", None)) self.thoughts.append(ans_dict.get("think", None)) - agent_info = AgentInfo( - think=ans_dict.get("think", None), - chat_messages=chat_messages, - stats=stats, - extra_info={"chat_model_args": asdict(self.chat_model_args)}, - ) + agent_info = make_agent_info(self.chat_llm, ans_dict, chat_messages, self.chat_model_args) return ans_dict["action"], agent_info def reset(self, seed=None): diff --git a/src/agentlab/agents/most_basic_agent/most_basic_agent.py b/src/agentlab/agents/most_basic_agent/most_basic_agent.py index 38145d5da..545c14f79 100644 --- a/src/agentlab/agents/most_basic_agent/most_basic_agent.py +++ b/src/agentlab/agents/most_basic_agent/most_basic_agent.py @@ -4,7 +4,7 @@ import bgym -from agentlab.agents.agent_args import AgentArgs +from agentlab.agents.agent_args import ChatModelAgentArgs from agentlab.experiments.loop import ExpArgs from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT from agentlab.llm.llm_utils import ( @@ -22,7 +22,7 @@ @dataclass -class MostBasicAgentArgs(AgentArgs): +class MostBasicAgentArgs(ChatModelAgentArgs): agent_name: str = "BasicAgent" temperature: float = 0.1 use_chain_of_thought: bool = False @@ -38,12 +38,6 @@ def make_agent(self) -> bgym.Agent: def set_reproducibility_mode(self): self.temperature = 0 - def prepare(self): - return self.chat_model_args.prepare_server() - - def close(self): - return self.chat_model_args.close_server() - class MostBasicAgent(bgym.Agent): def __init__( diff --git a/src/agentlab/agents/visual_agent/visual_agent.py b/src/agentlab/agents/visual_agent/visual_agent.py index d76cedf33..cd5df0133 100644 --- a/src/agentlab/agents/visual_agent/visual_agent.py +++ b/src/agentlab/agents/visual_agent/visual_agent.py @@ -8,14 +8,15 @@ the agent, including model arguments and flags for various behaviors. """ -from dataclasses import asdict, dataclass +from dataclasses import dataclass import bgym from bgym import Benchmark -from browsergym.experiments.agent import Agent, AgentInfo +from browsergym.experiments.agent import Agent from agentlab.agents import dynamic_prompting as dp -from agentlab.agents.agent_args import AgentArgs +from agentlab.agents.agent_args import ChatModelAgentArgs +from agentlab.agents.agent_utils import busted_retry_ans_dict, make_agent_info from agentlab.llm.chat_api import BaseModelArgs from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry from agentlab.llm.tracking import cost_tracker_decorator @@ -24,7 +25,7 @@ @dataclass -class VisualAgentArgs(AgentArgs): +class VisualAgentArgs(ChatModelAgentArgs): chat_model_args: BaseModelArgs = None flags: PromptFlags = None max_retry: int = 4 @@ -42,12 +43,6 @@ def set_benchmark(self, benchmark: Benchmark, demo_mode): def set_reproducibility_mode(self): self.chat_model_args.temperature = 0 - def prepare(self): - return self.chat_model_args.prepare_server() - - def close(self): - return self.chat_model_args.close_server() - def make_agent(self): return VisualAgent( chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry @@ -103,25 +98,12 @@ def get_action(self, obs): # inferring the number of retries, TODO: make this less hacky ans_dict["n_retry"] = (len(chat_messages) - 3) / 2 except ParseError: - ans_dict = dict( - action=None, - n_retry=self.max_retry + 1, - busted_retry=1, - ) - - stats = self.chat_llm.get_stats() - stats["n_retry"] = ans_dict["n_retry"] - stats["busted_retry"] = ans_dict["busted_retry"] + ans_dict = busted_retry_ans_dict(self.max_retry) self.actions.append(ans_dict["action"]) self.thoughts.append(ans_dict.get("think", None)) - agent_info = AgentInfo( - think=ans_dict.get("think", None), - chat_messages=chat_messages, - stats=stats, - extra_info={"chat_model_args": asdict(self.chat_model_args)}, - ) + agent_info = make_agent_info(self.chat_llm, ans_dict, chat_messages, self.chat_model_args) return ans_dict["action"], agent_info def reset(self, seed=None): diff --git a/src/agentlab/agents/visualwebarena/agent.py b/src/agentlab/agents/visualwebarena/agent.py index 465bd6aac..dc1dc7433 100644 --- a/src/agentlab/agents/visualwebarena/agent.py +++ b/src/agentlab/agents/visualwebarena/agent.py @@ -19,7 +19,7 @@ from agentlab.llm.llm_utils import ParseError, extract_code_blocks, retry from agentlab.llm.tracking import cost_tracker_decorator -from ..agent_args import AgentArgs +from ..agent_args import ChatModelAgentArgs from . import few_shots from .prompts import TEMPLATES @@ -50,7 +50,7 @@ def image_data_to_uri( @dataclass -class VisualWebArenaAgentArgs(AgentArgs): +class VisualWebArenaAgentArgs(ChatModelAgentArgs): agent_name: str = "VisualWebArenaAgent" temperature: float = 0.1 chat_model_args: BaseModelArgs = None @@ -80,12 +80,6 @@ def set_benchmark(self, benchmark: Benchmark, demo_mode: bool): def set_reproducibility_mode(self): self.temperature = 0.0 - def prepare(self): - return self.chat_model_args.prepare_server() - - def close(self): - return self.chat_model_args.close_server() - def parser(response: str) -> dict: blocks = extract_code_blocks(response) diff --git a/tests/agents/test_agent_utils.py b/tests/agents/test_agent_utils.py new file mode 100644 index 000000000..e4e09c167 --- /dev/null +++ b/tests/agents/test_agent_utils.py @@ -0,0 +1,86 @@ +from dataclasses import dataclass, field + +from browsergym.experiments.agent import AgentInfo + +from agentlab.agents.agent_args import ChatModelAgentArgs +from agentlab.agents.agent_utils import busted_retry_ans_dict, make_agent_info + + +def test_busted_retry_ans_dict(): + ans_dict = busted_retry_ans_dict(max_retry=4) + assert ans_dict == {"action": None, "n_retry": 5, "busted_retry": 1} + + +class FakeChatLLM: + def __init__(self, stats): + self._stats = stats + + def get_stats(self): + # return a fresh copy, mirroring real models that return per-call stats + return dict(self._stats) + + +@dataclass +class FakeChatModelArgs: + model_name: str = "fake/model" + temperature: float = 0.1 + + +def test_make_agent_info_collects_stats_and_packs_info(): + chat_llm = FakeChatLLM({"cost": 1.0}) + ans_dict = {"action": 'click("a1")', "think": "because", "n_retry": 2, "busted_retry": 0} + chat_messages = ["system", "human"] + chat_model_args = FakeChatModelArgs() + + info = make_agent_info(chat_llm, ans_dict, chat_messages, chat_model_args) + + assert isinstance(info, AgentInfo) + assert info.think == "because" + assert info.chat_messages == chat_messages + assert info.stats == {"cost": 1.0, "n_retry": 2, "busted_retry": 0} + assert info.extra_info == {"chat_model_args": {"model_name": "fake/model", "temperature": 0.1}} + + +def test_make_agent_info_defaults_think_to_none(): + info = make_agent_info( + FakeChatLLM({}), + {"action": None, "n_retry": 5, "busted_retry": 1}, + [], + FakeChatModelArgs(), + ) + assert info.think is None + + +class RecordingChatModelArgs: + """Minimal chat_model_args double recording prepare/close server calls.""" + + def __init__(self): + self.prepared = False + self.closed = False + + def prepare_server(self): + self.prepared = True + return "prepared" + + def close_server(self): + self.closed = True + return "closed" + + +@dataclass +class _DummyChatModelAgentArgs(ChatModelAgentArgs): + chat_model_args: object = None + + def make_agent(self): + raise NotImplementedError + + +def test_chat_model_agent_args_prepare_close_delegate(): + chat_model_args = RecordingChatModelArgs() + args = _DummyChatModelAgentArgs(chat_model_args=chat_model_args) + + assert args.prepare() == "prepared" + assert chat_model_args.prepared is True + + assert args.close() == "closed" + assert chat_model_args.closed is True