add support for tab visibility in observation flags and update related components

Author: recursix

src/agentlab/agents/dynamic_prompting.py

@@ -10,8 +10,6 @@
 
 import bgym
 from browsergym.core.action.base import AbstractActionSet
-from browsergym.core.action.highlevel import HighLevelActionSet
-from browsergym.core.action.python import PythonActionSet
 from browsergym.utils.obs import flatten_axtree_to_str, flatten_dom_to_str, overlay_som, prune_html
 
 from agentlab.llm.llm_utils import (
@@ -71,6 +69,7 @@ class ObsFlags(Flags):
 
     use_html: bool = True
     use_ax_tree: bool = False
+    use_tabs: bool = False
     use_focused_element: bool = False
     use_error_logs: bool = False
     use_history: bool = False
@@ -386,11 +385,7 @@ def _prompt(self) -> str:
     URL: {page_url}
 """
             prompt_pieces.append(prompt_piece)
-        self._prompt = "\n".join(prompt_pieces)
-
-
-def has_tab_action(action_set: bgym.HighLevelActionSetArgs):
-    return "tab" in action_set.subsets
+        return  "\n".join(prompt_pieces)
 
 
 class Observation(Shrinkable):
@@ -399,14 +394,14 @@ class Observation(Shrinkable):
     Contains the html, the accessibility tree and the error logs.
     """
 
-    def __init__(self, obs, flags: ObsFlags, use_tabs=False) -> None:
+    def __init__(self, obs, flags: ObsFlags) -> None:
         super().__init__()
         self.flags = flags
         self.obs = obs
 
         self.tabs = Tabs(
             obs,
-            visible=use_tabs,
+            visible=lambda: flags.use_tabs,
             prefix="## ",
         )
 

src/agentlab/agents/generic_agent/generic_agent.py

@@ -32,6 +32,7 @@ def set_benchmark(self, benchmark: bgym.Benchmark, demo_mode):
         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
@@ -268,5 +269,3 @@ def get_action_post_hoc(agent: GenericAgent, obs: dict, ans_dict: dict):
         output += f"\n<action>\n{action}\n</action>"
 
     return system_prompt, instruction_prompt, output
-    return system_prompt, instruction_prompt, output
-    return system_prompt, instruction_prompt, output

src/agentlab/agents/generic_agent/generic_agent_prompt.py

@@ -77,7 +77,6 @@ def __init__(
         self.obs = dp.Observation(
             obs_history[-1],
             self.flags.obs,
-            use_tabs=dp.has_tab_action(self.flags.action.action_set),
         )
 
         self.action_prompt = dp.ActionPrompt(action_set, action_flags=flags.action)

tests/agents/test_generic_prompt.py

@@ -20,33 +20,28 @@
 </html>
 """
 
+base_obs = {
+    "goal": "do this and that",
+    "goal_object": [{"type": "text", "text": "do this and that"}],
+    "chat_messages": [{"role": "user", "message": "do this and that"}],
+    "axtree_txt": "[1] Click me",
+    "focused_element_bid": "45-256",
+    "open_pages_urls": ["https://example.com"],
+    "open_pages_titles": ["Example"],
+    "active_page_index": 0,
+}
 
 OBS_HISTORY = [
-    {
-        "goal": "do this and that",
-        "goal_object": [{"type": "text", "text": "do this and that"}],
-        "chat_messages": [{"role": "user", "message": "do this and that"}],
+    base_obs | {
         "pruned_html": html_template.format(1),
-        "axtree_txt": "[1] Click me",
-        "focused_element_bid": "45-256",
         "last_action_error": "",
     },
-    {
-        "goal": "do this and that",
-        "goal_object": [{"type": "text", "text": "do this and that"}],
-        "chat_messages": [{"role": "user", "message": "do this and that"}],
+    base_obs | {
         "pruned_html": html_template.format(2),
-        "axtree_txt": "[1] Click me",
-        "focused_element_bid": "45-256",
         "last_action_error": "Hey, this is an error in the past",
     },
-    {
-        "goal": "do this and that",
-        "goal_object": [{"type": "text", "text": "do this and that"}],
-        "chat_messages": [{"role": "user", "message": "do this and that"}],
+    base_obs | {
         "pruned_html": html_template.format(3),
-        "axtree_txt": "[1] Click me",
-        "focused_element_bid": "45-256",
         "last_action_error": "Hey, there is an error now",
     },
 ]
@@ -58,6 +53,7 @@
     obs=dp.ObsFlags(
         use_html=True,
         use_ax_tree=True,
+        use_tabs=True,
         use_focused_element=True,
         use_error_logs=True,
         use_history=True,
@@ -104,6 +100,10 @@
         "obs.use_ax_tree",
         ("AXTree:", "Click me"),
     ),
+    (
+        "obs.use_tabs",
+        ("Currently open tabs:","(active tab)"),
+    ),
     (
         "obs.use_focused_element",
         ("Focused element:", "bid='45-256'"),
@@ -251,11 +251,6 @@ def test_main_prompt_elements_present():
     # for debugging
     test_shrinking_observation()
     test_main_prompt_elements_present()
-    for flag, expected_prompts in FLAG_EXPECTED_PROMPT:
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
-        test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
+    # for flag, expected_prompts in FLAG_EXPECTED_PROMPT:
+    #     test_main_prompt_elements_gone_one_at_a_time(flag, expected_prompts)
+