darglint and black

Author: amanjaiswal73892

src/agentlab/agents/hilt_agent/base_multi_candidate_agent.py

@@ -1,4 +1,5 @@
 from typing_extensions import Protocol
+
 from agentlab.agents.agent_args import AgentArgs
 
 
@@ -12,9 +13,10 @@ class MultiCandidateAgent(Protocol):
 
     def get_candidate_generations(
         self, obs: dict, hint: list[str] | None = None, n_candidates: int = 3
-    ) -> list[dict]:
+    ) -> "list[dict]":
         """
         Generate multiple candidate actions for the given observation.
+
         You can pass extra info in agent_info to update internal state of the
         agent based on the selected candidate. Your internal state management
         should be robust to multiple calls to the get_candidate_generations method
@@ -24,11 +26,6 @@ def get_candidate_generations(
             obs: The current observation dictionary containing environment state
             hint: Optional list of hint strings to guide candidate generation
             n_candidates: Number of candidate actions to generate
-
-        Returns:
-            List of dictionaries, each containing:
-                - 'action': The candidate action to be executed
-                - 'agent_info': Additional information about the action generation
         """
         ...
 
@@ -37,8 +34,10 @@ def update_agent_state_from_selected_candidate(self, output: dict):
         Update the agent's internal state based on the selected candidate.
         This can include any memory or planning updates.
 
+        Args:
+            output: The selected candidate action dictionary
         """
-        ...
+        pass
 
 
 class MultiCandidateAgentArgs(AgentArgs):
@@ -47,5 +46,5 @@ def make_agent(self) -> MultiCandidateAgent: ...
     def __post_init__(self):
         """Prefix subagent name with 'MC-'."""
         super().__post_init__()
-        if hasattr(self, 'agent_name') and self.agent_name:
+        if hasattr(self, "agent_name") and self.agent_name:
             self.agent_name = "MC-" + self.agent_name

src/agentlab/agents/hilt_agent/generic_human_guided_agent.py

@@ -7,6 +7,7 @@
 
 import bgym
 import numpy as np
+from browsergym.experiments.agent import AgentInfo
 from PIL import Image
 
 from agentlab.agents import dynamic_prompting as dp
@@ -23,7 +24,6 @@
     SystemMessage,
 )
 from agentlab.llm.tracking import cost_tracker_decorator
-from browsergym.experiments.agent import AgentInfo
 
 
 class CandidatesGeneration(dp.PromptElement):
@@ -87,15 +87,14 @@ def __init__(self, hint: list[str] | None = None, n_candidates=3) -> None:
     )
 
     def _parse_answer(self, text_answer: str) -> Dict[str, Dict[str, str]]:
-        """
-        Extract up to n_candidates candidates, using numbered tags only.
+        """Extract up to n_candidates candidates, using numbered tags only.
+
+        Args:
+            text_answer: The text response containing candidate generation tags.
 
         Returns:
-        {
-            "candidate_generation_1": {"think": "...", "action": "..."},
-            "candidate_generation_2": {"think": "...", "action": "..."},
-            ...
-        }
+            Dictionary mapping candidate names to their think and action content.
+            Format: {"candidate_generation_1": {"think": "...", "action": "..."}, ...}
         """
         result = {
             f"candidate_generation_{i+1}": {"think": "", "action": ""}
@@ -145,11 +144,11 @@ def make_agent(self):
         return MultipleProposalGenericAgent(
             chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
         )
-    
+
     def __post_init__(self):
         """Prefix subagent name with 'HILT-'."""
         super().__post_init__()
-        if hasattr(self, 'agent_name') and self.agent_name:
+        if hasattr(self, "agent_name") and self.agent_name:
             self.agent_name = "HILT-" + self.agent_name
 
 
@@ -363,13 +362,11 @@ def get_base_agent(llm_config):
     agent_configs = [HUMAN_GUIDED_GENERIC_AGENT]
     benchmark = bgym.DEFAULT_BENCHMARKS["miniwob"]()
     benchmark = benchmark.subset_from_glob("task_name", "*book*")
-    benchmark.env_args_list = benchmark.env_args_list[2:3]
+    benchmark.env_args_list = benchmark.env_args_list[3:4]
 
     for env_args in benchmark.env_args_list:
         env_args.max_steps = 100  # max human steps
-        env_args.headless = False
-        # env_args.use_chat_ui = False
-        # env_args.use_hint_labeling_ui = True
+        env_args.headless = True
 
     Study(agent_configs, benchmark, logging_level=logging.WARNING).run(
         n_jobs=1,

src/agentlab/agents/hilt_agent/hilt_agent.py

@@ -6,18 +6,19 @@
 
 import bgym
 import numpy as np
-from PIL import Image
 import playwright
+from browsergym.experiments.agent import Agent
+from PIL import Image
 
+from agentlab.agents.agent_args import AgentArgs
+from agentlab.agents.hilt_agent.base_multi_candidate_agent import MultiCandidateAgent
 from agentlab.agents.hilt_agent.hint_labelling import (
     HintLabeling,
     HintLabelingInputs,
 )
-from agentlab.llm.tracking import cost_tracker_decorator
 from agentlab.analyze import overlay_utils
-from browsergym.experiments.agent import Agent
-from agentlab.agents.agent_args import AgentArgs
-from agentlab.agents.hilt_agent.base_multi_candidate_agent import MultiCandidateAgent
+from agentlab.llm.tracking import cost_tracker_decorator
+
 
 class HumanInTheLoopAgent(Agent):
 
@@ -58,7 +59,7 @@ def get_action(self, obs):
         # Generate first candidates
         candidates = self.subagent.get_candidate_generations(obs, hint=None, n_candidates=3)
         step_n_human_intervention_rounds += 1
-        suggestions = [{ 'action': c['action'], 'think': c['agent_info'].think} for c in candidates]
+        suggestions = [{"action": c["action"], "think": c["agent_info"].think} for c in candidates]
         # List of Images as base64 - create overlay screenshots for each suggested action
         screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]
 
@@ -90,11 +91,11 @@ def get_action(self, obs):
                     hint = response["payload"]["hint"]
                     step_hint.append(hint)
                     candidates = self.subagent.get_candidate_generations(
-                        obs, 
-                        hint=step_hint if step_hint else None,
-                        n_candidates=3
+                        obs, hint=step_hint if step_hint else None, n_candidates=3
                     )
-                    suggestions = [{'action': c['action'], 'think': c['agent_info'].think} for c in candidates]
+                    suggestions = [
+                        {"action": c["action"], "think": c["agent_info"].think} for c in candidates
+                    ]
                     screenshots = [overlay_action(obs, choice["action"]) for choice in suggestions]
 
                 elif response["type"] == "step":
@@ -135,7 +136,6 @@ def get_action(self, obs):
 @dataclass
 class HumanInTheLoopAgentArgs(AgentArgs):
     subagent_args: Optional[AgentArgs] = None  # args for the underlying multiple proposal agent
-    
 
     def make_agent(self):
         assert self.subagent_args is not None
@@ -146,15 +146,15 @@ def __post_init__(self):
         super().__post_init__()
         if self.subagent_args and self.subagent_args.agent_name:
             self.agent_name = "HILT-" + self.subagent_args.agent_name
-    
+
     def set_benchmark(self, benchmark, demo_mode):
         """Delegate set_benchmark to the subagent if it has the method."""
-        if hasattr(self.subagent_args, 'set_benchmark'):
+        if hasattr(self.subagent_args, "set_benchmark"):
             self.subagent_args.set_benchmark(benchmark, demo_mode)
-    
+
     def set_reproducibility_mode(self):
         """Delegate set_reproducibility_mode to the subagent if it has the method."""
-        if hasattr(self.subagent_args, 'set_reproducibility_mode'):
+        if hasattr(self.subagent_args, "set_reproducibility_mode"):
             self.subagent_args.set_reproducibility_mode()
 
 
@@ -175,16 +175,17 @@ def img_to_base_64(image: Image.Image | np.ndarray) -> str:
     b64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
     return b64_str
 
+
 def get_base_human_in_the_loop_genericagent(llm_config):
     from agentlab.agents.generic_agent.tmlr_config import BASE_FLAGS
-    from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
     from agentlab.agents.hilt_agent.hilt_agent import HumanInTheLoopAgentArgs
     from agentlab.agents.hilt_agent.multi_candidate_generic_agent import (
         MultiCandidateGenericAgentArgs,
     )
+    from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
 
     return HumanInTheLoopAgentArgs(
-        subagent_args = MultiCandidateGenericAgentArgs(
+        subagent_args=MultiCandidateGenericAgentArgs(
             chat_model_args=CHAT_MODEL_ARGS_DICT[llm_config],
             flags=BASE_FLAGS,
         )
@@ -210,7 +211,6 @@ def get_base_human_in_the_loop_genericagent(llm_config):
         env_args.max_steps = 100  # max human steps
         env_args.headless = False
 
-
     Study(agent_configs, benchmark, logging_level=logging.WARNING).run(
         n_jobs=1,
         parallel_backend="sequential",

src/agentlab/agents/hilt_agent/hint_labelling.py

@@ -5,10 +5,10 @@
 from typing import Dict, List, Optional
 
 import playwright.sync_api
+from browsergym.core import _get_global_playwright
 from pydantic import BaseModel, Field
 
 from agentlab.agents.hilt_agent import hint_labelling_ui_files
-from browsergym.core import _get_global_playwright
 
 logger = logging.getLogger(__name__)
 
@@ -109,6 +109,16 @@ def wait_for_response(self, timeout: Optional[float] = 600) -> dict:
         """
         Wait until the page makes a request to /api/reprompt or /api/submit,
         then parse the request body and return it in your schema.
+
+        Args:
+            timeout (Optional[float]): Maximum time to wait for the request in seconds. If None or 0,
+                waits indefinitely. Defaults to 600 seconds.
+
+        Returns:
+            dict: A dictionary containing the parsed response with 'type' and 'payload' keys.
+                For /api/reprompt: {'type': 'reprompt', 'payload': {'hint': str}}
+                For /api/submit: {'type': 'step', 'payload': {'think': str, 'action': str}}
+
         """
         logger.info("Waiting for response from Hint Labeling UI...")
 

src/agentlab/agents/hilt_agent/multi_candidate_generic_agent.py

@@ -2,11 +2,12 @@
 from dataclasses import asdict, dataclass
 from typing import Dict, List
 
+from browsergym.experiments.agent import AgentInfo
+
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.agents.generic_agent.generic_agent import GenericAgent, GenericAgentArgs
 from agentlab.agents.generic_agent.generic_agent_prompt import MainPrompt
 from agentlab.llm.llm_utils import Discussion, HumanMessage, SystemMessage
-from browsergym.experiments.agent import AgentInfo
 
 
 class CandidatesGeneration(dp.PromptElement):
@@ -70,15 +71,14 @@ def __init__(self, hint: list[str] | None = None, n_candidates=3) -> None:
     )
 
     def _parse_answer(self, text_answer: str) -> Dict[str, Dict[str, str]]:
-        """
-        Extract up to n_candidates candidates, using numbered tags only.
+        """Extract up to n_candidates candidates, using numbered tags only.
+
+        Args:
+            text_answer: The text response containing candidate generation tags.
 
         Returns:
-        {
-            "candidate_generation_1": {"think": "...", "action": "..."},
-            "candidate_generation_2": {"think": "...", "action": "..."},
-            ...
-        }
+            Dictionary mapping candidate names to their think and action content.
+            Format: {"candidate_generation_1": {"think": "...", "action": "..."}, ...}
         """
         result = {
             f"candidate_generation_{i+1}": {"think": "", "action": ""}
@@ -123,7 +123,6 @@ def get_candidate_generations(
         # Important to handle cases when get_candidate_generation is called multiple times in a single step.
         if not self.obs_history or self.obs_history[-1] is not obs:
             self.obs_history.append(obs)
-    
 
         main_prompt = MainPrompt(
             action_set=self.action_set,
@@ -180,8 +179,12 @@ def get_candidate_generations(
         return output
 
     def update_agent_state_from_selected_candidate(self, output):
-        """Updates the agent's internal state based on the selected candidate from human feedback."""
-        action, agent_info = output['action'], output['agent_info']
+        """Updates the agent's internal state based on the selected candidate from human feedback.
+
+        Args:
+            output: Dictionary containing 'action' and 'agent_info' keys from selected candidate.
+        """
+        action, agent_info = output["action"], output["agent_info"]
         self.plan = agent_info.extra_info.get("plan", self.plan)
         self.plan_step = agent_info.extra_info.get("step", self.plan_step)
         self.memories.append(agent_info.extra_info.get("memory", None))
@@ -191,11 +194,17 @@ def update_agent_state_from_selected_candidate(self, output):
     def get_action(self, obs):
         """Generates multiple candidates and always returns the first one.
         This allows to use this agent as a drop-in replacement for a single-candidate agent.
+
+        Args:
+            obs: The observation from the environment.
+
+        Returns:
+            tuple: A tuple containing (action, agent_info).
         """
-        candidates = self.get_candidate_generations(obs, hint=None, n_candidates=2) 
-        selection = candidates[0] # always select the first option.
+        candidates = self.get_candidate_generations(obs, hint=None, n_candidates=2)
+        selection = candidates[0]  # always select the first option.
         self.update_agent_state_from_selected_candidate(selection)
-        action, agent_info = selection['action'], selection['agent_info']
+        action, agent_info = selection["action"], selection["agent_info"]
 
         return action, agent_info
 
@@ -212,5 +221,5 @@ def make_agent(self):
     def __post_init__(self):
         """Prefix subagent name with 'MC-'."""
         super().__post_init__()
-        if hasattr(self, 'agent_name') and self.agent_name:
+        if hasattr(self, "agent_name") and self.agent_name:
             self.agent_name = "MC-" + self.agent_name