Enhance StructuredDiscussion to group messages with summaries and adjust ToolUseAgent configuration for multi-action support

Author: recursix

src/agentlab/agents/tool_use_agent/multi_tool_agent.py

@@ -1,4 +1,5 @@
 import fnmatch
+import json
 import logging
 from abc import ABC, abstractmethod
 from copy import copy
@@ -62,6 +63,12 @@ class MsgGroup:
 
 
 class StructuredDiscussion:
+    """
+    A structured discussion that groups messages into named groups with a potential summary for each group.
+
+    When the discussion is flattened, only the last `keep_last_n_obs` groups are kept in the final list,
+    the other groups are replaced by their summaries if they have one.
+    """
 
     def __init__(self, keep_last_n_obs=None):
         self.groups: list[MsgGroup] = []
@@ -84,9 +91,7 @@ def flatten(self) -> list[MessageBuilder]:
         messages = []
         for i, group in enumerate(self.groups):
             is_tail = i >= len(self.groups) - keep_last_n_obs
-            print(
-                f"Processing group {i} ({group.name}), is_tail={is_tail}, len(greoup)={len(group.messages)}"
-            )
+
             if not is_tail and group.summary is not None:
                 messages.append(group.summary)
             else:
@@ -103,15 +108,6 @@ def is_goal_set(self) -> bool:
         return len(self.groups) > 0
 
 
-# @dataclass
-# class BlockArgs(ABC):
-
-#     @abstractmethod
-#     def make(self) -> Block:
-#         """Make a block from the arguments."""
-#         return self.__class__(**asdict(self))
-
-
 SYS_MSG = """You are a web agent. Based on the observation, you will decide which action to take to accomplish your goal. 
 You strive for excellence and need to be as meticulous as possible. Make sure to explore when not sure.
 """
@@ -344,7 +340,9 @@ class PromptConfig:
     summarizer: Summarizer = None
     general_hints: GeneralHints = None
     task_hint: TaskHint = None
-    keep_last_n_obs: int = 2
+    keep_last_n_obs: int = 1
+    multiaction: bool = False
+    action_subsets: tuple[str] = field(default_factory=lambda: ("coord",))
 
 
 @dataclass
@@ -382,7 +380,9 @@ def __init__(
     ):
         self.model_args = model_args
         self.config = config
-        self.action_set = bgym.HighLevelActionSet(["coord"], multiaction=False)
+        self.action_set = bgym.HighLevelActionSet(
+            self.config.action_subsets, multiaction=self.config.multiaction
+        )
         self.tools = self.action_set.to_tool_description(api=model_args.api)
 
         self.call_ids = []
@@ -444,7 +444,6 @@ def get_action(self, obs: Any) -> float:
             self.discussion.new_group()
 
         self.obs_block.apply(self.llm, self.discussion, obs, last_llm_output=self.last_response)
-        print("flatten for summary")
 
         self.config.summarizer.apply(self.llm, self.discussion)
 
@@ -466,6 +465,13 @@ def get_action(self, obs: Any) -> float:
         self._responses.append(response)  # may be useful for debugging
         # self.messages.append(response.assistant_message)  # this is tool call
 
+        tools_str = json.dumps(self.tools, indent=2)
+        tools_msg = MessageBuilder("tool_description").add_text(tools_str)
+
+        # Adding these extra messages to visualize in gradio
+        messages.insert(0, tools_msg)  # insert at the beginning of the messages
+        messages.append(response.tool_calls)
+
         agent_info = bgym.AgentInfo(
             think=think,
             chat_messages=messages,
@@ -509,13 +515,16 @@ def get_action(self, obs: Any) -> float:
         use_last_error=True,
         use_screenshot=True,
         use_axtree=False,
-        use_dom=False,
+        use_dom=True,
         use_som=False,
         use_tabs=False,
     ),
     summarizer=Summarizer(do_summary=True),
     general_hints=GeneralHints(use_hints=False),
     task_hint=TaskHint(use_task_hint=True),
+    keep_last_n_obs=1,  # keep only the last observation in the discussion
+    multiaction=False,  # whether to use multi-action or not
+    action_subsets=("bid",),
 )
 
 AGENT_CONFIG = ToolUseAgentArgs(