Implement VisualAgent and associated prompt flags for enhanced agent functionality

Author: recursix

src/agentlab/agents/visual_agent/agent_configs.py

@@ -0,0 +1,48 @@
+from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
+
+from .visual_agent import VisualAgentArgs
+from .visual_agent_prompts import PromptFlags
+import agentlab.agents.dynamic_prompting as dp
+import bgym
+
+# the other flags are ignored for this agent.
+DEFAULT_OBS_FLAGS = dp.ObsFlags(
+    use_tabs=True,  # will be overridden by the benchmark when set_benchmark is called after initalizing the agent
+    use_error_logs=True,
+    use_past_error_logs=False,
+    use_screenshot=True,
+    use_som=False,
+    openai_vision_detail="auto",
+)
+
+DEFAULT_ACTION_FLAGS = dp.ActionFlags(
+    action_set=bgym.HighLevelActionSetArgs(subsets=["coord"]),
+    long_description=True,
+    individual_examples=False,
+)
+
+
+DEFAULT_PROMPT_FLAGS = PromptFlags(
+    obs=DEFAULT_OBS_FLAGS,
+    action=DEFAULT_ACTION_FLAGS,
+    use_thinking=True,
+    use_concrete_example=False,
+    use_abstract_example=True,
+    enable_chat=False,
+    extra_instructions=None,
+)
+
+VISUAL_AGENT_4o = VisualAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["openai/gpt-4o-2024-05-13"],
+    flags=DEFAULT_PROMPT_FLAGS,
+)
+
+VISUAL_AGENT_COMPUTER_USE = VisualAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["openai/computer-use-preview-2025-03-11"],
+    flags=DEFAULT_PROMPT_FLAGS,
+)
+
+VISUAL_AGENT_CLAUDE_3_5 = VisualAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["openrouter/anthropic/claude-3.5-sonnet:beta"],
+    flags=DEFAULT_PROMPT_FLAGS,
+)

src/agentlab/agents/visual_agent/visual_agent.py

@@ -8,9 +8,7 @@
 the agent, including model arguments and flags for various behaviors.
 """
 
-from copy import deepcopy
 from dataclasses import asdict, dataclass
-from warnings import warn
 
 import bgym
 from browsergym.experiments.agent import Agent, AgentInfo
@@ -21,44 +19,24 @@
 from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
 from agentlab.llm.tracking import cost_tracker_decorator
 
-from .visual_agent_prompts import GenericPromptFlags, MainPrompt
-from functools import partial
+from .visual_agent_prompts import PromptFlags, MainPrompt
 
 
 @dataclass
-class ToolAgentFlags:
-    pass
-
-
-@dataclass
-class ToolAgentArgs(AgentArgs):
+class VisualAgentArgs(AgentArgs):
     chat_model_args: BaseModelArgs = None
-    flags: GenericPromptFlags = None
+    flags: PromptFlags = None
     max_retry: int = 4
 
     def __post_init__(self):
-        try:  # some attributes might be temporarily args.CrossProd for hyperparameter generation
-            self.agent_name = f"GenericAgent-{self.chat_model_args.model_name}".replace("/", "_")
+        try:  # some attributes might be missing temporarily due to args.CrossProd for hyperparameter generation
+            self.agent_name = f"VisualAgent-{self.chat_model_args.model_name}".replace("/", "_")
         except AttributeError:
             pass
 
     def set_benchmark(self, benchmark: bgym.Benchmark, demo_mode):
         """Override Some flags based on the benchmark."""
-        if benchmark.name.startswith("miniwob"):
-            self.flags.obs.use_html = True
-
         self.flags.obs.use_tabs = benchmark.is_multi_tab
-        self.flags.action.action_set = deepcopy(benchmark.high_level_action_set_args)
-
-        # for backward compatibility with old traces
-        if self.flags.action.multi_actions is not None:
-            self.flags.action.action_set.multiaction = self.flags.action.multi_actions
-        if self.flags.action.is_strict is not None:
-            self.flags.action.action_set.strict = self.flags.action.is_strict
-
-        # verify if we can remove this
-        if demo_mode:
-            self.flags.action.action_set.demo_mode = "all_blue"
 
     def set_reproducibility_mode(self):
         self.chat_model_args.temperature = 0
@@ -70,17 +48,17 @@ def close(self):
         return self.chat_model_args.close_server()
 
     def make_agent(self):
-        return ToolAgent(
+        return VisualAgent(
             chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
         )
 
 
-class ToolAgent(Agent):
+class VisualAgent(Agent):
 
     def __init__(
         self,
         chat_model_args: BaseModelArgs,
-        flags: GenericPromptFlags,
+        flags: PromptFlags,
         max_retry: int = 4,
     ):
 
@@ -92,7 +70,6 @@ def __init__(
         self.action_set = self.flags.action.action_set.make_action_set()
         self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)
 
-        self._check_flag_constancy()
         self.reset(seed=None)
 
     def obs_preprocessor(self, obs: dict) -> dict:
@@ -101,34 +78,20 @@ def obs_preprocessor(self, obs: dict) -> dict:
     @cost_tracker_decorator
     def get_action(self, obs):
 
-        self.obs_history.append(obs)
         main_prompt = MainPrompt(
             action_set=self.action_set,
-            obs_history=self.obs_history,
+            obs=obs,
             actions=self.actions,
-            memories=self.memories,
             thoughts=self.thoughts,
-            previous_plan=self.plan,
-            step=self.plan_step,
             flags=self.flags,
         )
 
-        max_prompt_tokens, max_trunc_itr = self._get_maxes()
-
         system_prompt = SystemMessage(dp.SystemPrompt().prompt)
-
-        human_prompt = dp.fit_tokens(
-            shrinkable=main_prompt,
-            max_prompt_tokens=max_prompt_tokens,
-            model_name=self.chat_model_args.model_name,
-            max_iterations=max_trunc_itr,
-            additional_prompts=system_prompt,
-        )
         try:
             # TODO, we would need to further shrink the prompt if the retry
             # cause it to be too long
 
-            chat_messages = Discussion([system_prompt, human_prompt])
+            chat_messages = Discussion([system_prompt, main_prompt.prompt])
             ans_dict = retry(
                 self.chat_llm,
                 chat_messages,
@@ -138,7 +101,7 @@ 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:
+        except ParseError:
             ans_dict = dict(
                 action=None,
                 n_retry=self.max_retry + 1,
@@ -149,10 +112,7 @@ def get_action(self, obs):
         stats["n_retry"] = ans_dict["n_retry"]
         stats["busted_retry"] = ans_dict["busted_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(
@@ -165,43 +125,5 @@ def get_action(self, obs):
 
     def reset(self, seed=None):
         self.seed = seed
-        self.plan = "No plan yet"
-        self.plan_step = -1
-        self.memories = []
         self.thoughts = []
         self.actions = []
-        self.obs_history = []
-
-    def _check_flag_constancy(self):
-        flags = self.flags
-        if flags.obs.use_som:
-            if not flags.obs.use_screenshot:
-                warn(
-                    """
-Warning: use_som=True requires use_screenshot=True. Disabling use_som."""
-                )
-                flags.obs.use_som = False
-        if flags.obs.use_screenshot:
-            if not self.chat_model_args.vision_support:
-                warn(
-                    """
-Warning: use_screenshot is set to True, but the chat model \
-does not support vision. Disabling use_screenshot."""
-                )
-                flags.obs.use_screenshot = False
-        return flags
-
-    def _get_maxes(self):
-        maxes = (
-            self.flags.max_prompt_tokens,
-            self.chat_model_args.max_total_tokens,
-            self.chat_model_args.max_input_tokens,
-        )
-        maxes = [m for m in maxes if m is not None]
-        max_prompt_tokens = min(maxes) if maxes else None
-        max_trunc_itr = (
-            self.flags.max_trunc_itr
-            if self.flags.max_trunc_itr
-            else 20  # dangerous to change the default value here?
-        )
-        return max_prompt_tokens, max_trunc_itr

src/agentlab/agents/visual_agent/visual_agent_prompts.py

@@ -6,6 +6,7 @@
 
 import logging
 from dataclasses import dataclass
+import bgym
 
 from browsergym.core.action.base import AbstractActionSet
 
@@ -17,35 +18,15 @@
 class PromptFlags(dp.Flags):
     """
     A class to represent various flags used to control features in an application.
-
-    Attributes:
-        use_criticise (bool): Ask the LLM to first draft and criticise the action before producing it.
-        use_thinking (bool): Enable a chain of thoughts.
-        use_concrete_example (bool): Use a concrete example of the answer in the prompt for a generic task.
-        use_abstract_example (bool): Use an abstract example of the answer in the prompt.
-        use_hints (bool): Add some human-engineered hints to the prompt.
-        enable_chat (bool): Enable chat mode, where the agent can interact with the user.
-        max_prompt_tokens (int): Maximum number of tokens allowed in the prompt.
-        be_cautious (bool): Instruct the agent to be cautious about its actions.
-        extra_instructions (Optional[str]): Extra instructions to provide to the agent.
-        add_missparsed_messages (bool): When retrying, add the missparsed messages to the prompt.
-        flag_group (Optional[str]): Group of flags used.
     """
 
-    obs: dp.ObsFlags
-    action: dp.ActionFlags
-    use_criticise: bool = False  #
-    use_thinking: bool = False
-    use_concrete_example: bool = True
-    use_abstract_example: bool = False
-    use_hints: bool = False
+    obs: dp.ObsFlags = None
+    action: dp.ActionFlags = None
+    use_thinking: bool = True
+    use_concrete_example: bool = False
+    use_abstract_example: bool = True
     enable_chat: bool = False
-    max_prompt_tokens: int = None
-    be_cautious: bool = True
     extra_instructions: str | None = None
-    add_missparsed_messages: bool = True
-    max_trunc_itr: int = 20
-    flag_group: str = None
 
 
 class SystemPrompt(dp.PromptElement):
@@ -77,7 +58,7 @@ class History(dp.PromptElement):
     Format the actions and thoughts of previous steps."""
 
     def __init__(self, actions, thoughts) -> None:
-
+        super().__init__()
         prompt_elements = []
         for i, (action, thought) in enumerate(zip(actions, thoughts)):
             prompt_elements.append(
@@ -121,7 +102,7 @@ def __init__(self, obs, flags: dp.ObsFlags) -> None:
     def _prompt(self) -> str:
         return f"""
 # Observation of current step:
-{self.tabs.prompt}{self.focused_element.prompt}{self.error.prompt}
+{self.tabs.prompt}{self.error.prompt}
 
 """
 
@@ -152,12 +133,9 @@ def __init__(
     ) -> None:
         super().__init__()
         self.flags = flags
-        self.history = History(obs, actions, thoughts)
+        self.history = History(actions, thoughts)
         self.instructions = make_instructions(obs, flags.enable_chat, flags.extra_instructions)
-        self.obs = dp.Observation(
-            obs,
-            self.flags.obs,
-        )
+        self.obs = Observation(obs, self.flags.obs)
 
         self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)
         self.think = dp.Think(visible=lambda: flags.use_thinking)