feature: improve agent code reusability by extracting shared boilerplate

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()

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.

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,29 +129,16 @@ 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)
         self.actions.append(ans_dict["action"])
         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):

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__(

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):

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)

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