fix tests

Author: recursix

src/agentlab/agents/dynamic_prompting.py

@@ -10,7 +10,12 @@
 
 import bgym
 from browsergym.core.action.base import AbstractActionSet
-from browsergym.utils.obs import flatten_axtree_to_str, flatten_dom_to_str, overlay_som, prune_html
+from browsergym.utils.obs import (
+    flatten_axtree_to_str,
+    flatten_dom_to_str,
+    overlay_som,
+    prune_html,
+)
 
 from agentlab.llm.llm_utils import (
     BaseMessage,
@@ -385,7 +390,7 @@ def _prompt(self) -> str:
     URL: {page_url}
 """
             prompt_pieces.append(prompt_piece)
-        return  "\n".join(prompt_pieces)
+        return "\n".join(prompt_pieces)
 
 
 class Observation(Shrinkable):

src/agentlab/experiments/graph_execution_ray.py

@@ -2,12 +2,14 @@
 
 # # Disable Ray log deduplication
 # os.environ["RAY_DEDUP_LOGS"] = "0"
+import logging
 import time
-import ray
+
 import bgym
-from agentlab.experiments.exp_utils import run_exp, _episode_timeout
+import ray
 from ray.util import state
-import logging
+
+from agentlab.experiments.exp_utils import _episode_timeout, run_exp
 
 logger = logging.getLogger(__name__)
 
@@ -36,7 +38,7 @@ def get_task(exp_arg: bgym.ExpArgs):
         get_task(exp_arg)
 
     max_timeout = max([_episode_timeout(exp_args, avg_step_timeout) for exp_args in exp_args_list])
-    
+
     return poll_for_timeout(task_map, max_timeout, poll_interval=max_timeout * 0.1)
 
 

src/agentlab/experiments/study.py

@@ -254,7 +254,7 @@ def get_report(self, ignore_cache=False, ignore_stale=False):
         return inspect_results.get_study_summary(
             self.dir, ignore_cache=ignore_cache, ignore_stale=ignore_stale
         )
-    
+
     def override_max_steps(self, max_steps):
         for exp_args in self.exp_args_list:
             exp_args.env_args.max_steps = max_steps

tests/agents/test_generic_prompt.py

@@ -5,7 +5,10 @@
 
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_3_5
-from agentlab.agents.generic_agent.generic_agent_prompt import GenericPromptFlags, MainPrompt
+from agentlab.agents.generic_agent.generic_agent_prompt import (
+    GenericPromptFlags,
+    MainPrompt,
+)
 from agentlab.llm.llm_utils import count_tokens
 
 html_template = """
@@ -32,15 +35,18 @@
 }
 
 OBS_HISTORY = [
-    base_obs | {
+    base_obs
+    | {
         "pruned_html": html_template.format(1),
         "last_action_error": "",
     },
-    base_obs | {
+    base_obs
+    | {
         "pruned_html": html_template.format(2),
         "last_action_error": "Hey, this is an error in the past",
     },
-    base_obs | {
+    base_obs
+    | {
         "pruned_html": html_template.format(3),
         "last_action_error": "Hey, there is an error now",
     },
@@ -102,7 +108,7 @@
     ),
     (
         "obs.use_tabs",
-        ("Currently open tabs:","(active tab)"),
+        ("Currently open tabs:", "(active tab)"),
     ),
     (
         "obs.use_focused_element",
@@ -165,7 +171,7 @@ def test_shrinking_observation():
     flags.obs.use_html = True
 
     prompt_maker = MainPrompt(
-        action_set=dp.HighLevelActionSet(),
+        action_set=bgym.HighLevelActionSet(),
         obs_history=OBS_HISTORY,
         actions=ACTIONS,
         memories=MEMORIES,
@@ -231,7 +237,7 @@ def test_main_prompt_elements_present():
     # Initialize MainPrompt
     prompt = str(
         MainPrompt(
-            action_set=dp.HighLevelActionSet(),
+            action_set=bgym.HighLevelActionSet(),
             obs_history=OBS_HISTORY,
             actions=ACTIONS,
             memories=MEMORIES,
@@ -253,4 +259,3 @@ def test_main_prompt_elements_present():
     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)
-