Merge branch 'generic_agent_hinter' into step-wise-retrieval

Author: hnekoeiq

pyproject.toml

@@ -109,3 +109,4 @@ hint = [
 [project.scripts]
 agentlab-assistant = "agentlab.ui_assistant:main"
 agentlab-xray = "agentlab.analyze.agent_xray:main"
+agentlab-mentor = "agentlab.agents.hitl_agent.launch_hint_ui:main"

src/agentlab/agents/agent_utils.py

@@ -1,6 +1,11 @@
+import copy
+
 from PIL import Image, ImageDraw
 from playwright.sync_api import Page
 
+from agentlab.analyze import overlay_utils
+from agentlab.llm.llm_utils import img_to_base_64
+
 
 def draw_mouse_pointer(image: Image.Image, x: int, y: int) -> Image.Image:
     """
@@ -128,3 +133,24 @@ def zoom_webpage(page: Page, zoom_factor: float = 1.5):
 
     page.evaluate(f"document.documentElement.style.zoom='{zoom_factor*100}%'")
     return page
+
+
+def overlay_action(obs, action):
+    """Overlays actions on screenshot in-place"""
+    act_img = copy.deepcopy(obs["screenshot"])
+    act_img = Image.fromarray(act_img)
+
+    new_obs_properties = copy.deepcopy(obs["extra_element_properties"])
+    import os
+
+    if os.getenv("AGENTLAB_USE_RETINA"):
+        # HACK: divide everything by 2 in the obs
+        # TODO: make this more robust by changing login in annotate_action directly (or maybe in the obs section?)
+        for key, value in new_obs_properties.items():
+            try:
+                new_obs_properties[key]["bbox"] = [elem / 2 for elem in value["bbox"]]
+            except:
+                pass
+
+    overlay_utils.annotate_action(act_img, action, properties=new_obs_properties)
+    return img_to_base_64(act_img)

src/agentlab/agents/generic_agent/__init__.py

@@ -22,6 +22,22 @@
     AGENT_4o_MINI,
     AGENT_4o_MINI_VISION,
     AGENT_4o_VISION,
+    AGENT_AZURE_4o_MINI,
+    AGENT_AZURE_4o,
+    AGENT_AZURE_4o_VISION,
+    AGENT_AZURE_4o_MINI_VISION,
+    AGENT_AZURE_41,
+    AGENT_AZURE_41_MINI,
+    AGENT_AZURE_41_NANO,
+    AGENT_AZURE_41_VISION,
+    AGENT_AZURE_41_MINI_VISION,
+    AGENT_AZURE_41_NANO_VISION,
+    AGENT_AZURE_5,
+    AGENT_AZURE_5_MINI,
+    AGENT_AZURE_5_NANO,
+    AGENT_AZURE_5_VISION,
+    AGENT_AZURE_5_MINI_VISION,
+    AGENT_AZURE_5_NANO_VISION,
     AGENT_o1_MINI,
     AGENT_o3_MINI,
     FLAGS_GPT_4o,
@@ -46,6 +62,22 @@
     "AGENT_37_SONNET",
     "AGENT_4o_VISION",
     "AGENT_4o_MINI_VISION",
+    "AGENT_AZURE_4o_MINI",
+    "AGENT_AZURE_4o",
+    "AGENT_AZURE_4o_VISION",
+    "AGENT_AZURE_4o_MINI_VISION",
+    "AGENT_AZURE_41",
+    "AGENT_AZURE_41_MINI",
+    "AGENT_AZURE_41_NANO",
+    "AGENT_AZURE_41_VISION",
+    "AGENT_AZURE_41_MINI_VISION",
+    "AGENT_AZURE_41_NANO_VISION",
+    "AGENT_AZURE_5",
+    "AGENT_AZURE_5_MINI",
+    "AGENT_AZURE_5_NANO",
+    "AGENT_AZURE_5_VISION",
+    "AGENT_AZURE_5_MINI_VISION",
+    "AGENT_AZURE_5_NANO_VISION",
     "AGENT_CLAUDE_SONNET_35_VISION",
     "AGENT_GPT5_MINI",
 ]

src/agentlab/agents/generic_agent/agent_configs.py

@@ -262,6 +262,43 @@
     chat_model_args=CHAT_MODEL_ARGS_DICT["openai/gpt-4o-mini-2024-07-18"],
     flags=FLAGS_GPT_4o,
 )
+
+AGENT_AZURE_4o_MINI = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4o-mini-2024-07-18"],
+    flags=FLAGS_GPT_4o,
+)
+AGENT_AZURE_4o = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4o-2024-08-06"],
+    flags=FLAGS_GPT_4o,
+)
+AGENT_AZURE_41 = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-2025-04-14"],
+    flags=FLAGS_GPT_4o,
+)
+AGENT_AZURE_41_MINI = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-mini-2025-04-14"],
+    flags=FLAGS_GPT_4o,
+)
+AGENT_AZURE_41_NANO = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-nano-2025-04-14"],
+    flags=FLAGS_GPT_4o,
+)
+
+AGENT_AZURE_5 = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-2025-08-07"],
+    flags=FLAGS_GPT_4o,
+)
+
+AGENT_AZURE_5_MINI = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-mini-2025-08-07"],
+    flags=FLAGS_GPT_4o,
+)
+
+AGENT_AZURE_5_NANO = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-nano-2025-08-07"],
+    flags=FLAGS_GPT_4o,
+)
+
 AGENT_CLAUDE_SONNET_35 = GenericAgentArgs(
     chat_model_args=CHAT_MODEL_ARGS_DICT["openrouter/anthropic/claude-3.5-sonnet:beta"],
     flags=FLAGS_GPT_4o,
@@ -298,6 +335,45 @@
     flags=FLAGS_GPT_4o_VISION,
 )
 
+AGENT_AZURE_4o_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4o-2024-08-06"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
+AGENT_AZURE_4o_MINI_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4o-mini-2024-07-18"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
+AGENT_AZURE_41_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-2025-04-14"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
+AGENT_AZURE_41_MINI_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-mini-2025-04-14"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+AGENT_AZURE_41_NANO_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-nano-2025-04-14"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
+AGENT_AZURE_5_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-2025-08-07"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
+AGENT_AZURE_5_MINI_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-mini-2025-08-07"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
+AGENT_AZURE_5_NANO_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-nano-2025-08-07"],
+    flags=FLAGS_GPT_4o_VISION,
+)
+
 AGENT_CLAUDE_SONNET_35_VISION = GenericAgentArgs(
     chat_model_args=CHAT_MODEL_ARGS_DICT["openrouter/anthropic/claude-3.5-sonnet:beta"],
     flags=FLAGS_GPT_4o_VISION,

src/agentlab/agents/generic_agent_hinter/generic_agent.py

@@ -38,7 +38,7 @@ class GenericAgentArgs(AgentArgs):
 
     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("/", "_")
+            self.agent_name = f"GenericAgent-hinter-{self.chat_model_args.model_name}".replace("/", "_")
         except AttributeError:
             pass
 

src/agentlab/agents/generic_agent_hinter/generic_agent_prompt.py

@@ -61,6 +61,13 @@ class GenericPromptFlags(dp.Flags):
     add_missparsed_messages: bool = True
     max_trunc_itr: int = 20
     flag_group: str = None
+    # hint flags
+    hint_type: Literal["human", "llm", "docs"] = "human"
+    hint_index_type: Literal["sparse", "dense"] = "sparse"
+    hint_query_type: Literal["direct", "llm", "emb"] = "direct"
+    hint_index_path: str = None
+    hint_retriever_path: str = None
+    hint_num_results: int = 5
     n_retrieval_queries: int = 3
     hint_level: Literal["episode", "step"] = "episode"
 
@@ -120,6 +127,13 @@ def time_for_caution():
             hint_retrieval_mode=flags.task_hint_retrieval_mode,
             llm=llm,
             skip_hints_for_current_task=flags.skip_hints_for_current_task,
+            # hint related
+            hint_type=flags.hint_type,
+            hint_index_type=flags.hint_index_type,
+            hint_query_type=flags.hint_query_type,
+            hint_index_path=flags.hint_index_path,
+            hint_retriever_path=flags.hint_retriever_path,
+            hint_num_results=flags.hint_num_results,
             hint_level=flags.hint_level,
             queries=queries,
         )
@@ -307,14 +321,26 @@ def __init__(
         use_task_hint: bool,
         hint_db_path: str,
         goal: str,
-        hint_retrieval_mode: Literal["direct", "llm", "emb"],
-        skip_hints_for_current_task: bool,
         llm: ChatModel,
+        hint_type: Literal["human", "llm", "docs"] = "human",
+        hint_index_type: Literal["sparse", "dense"] = "sparse",
+        hint_query_type: Literal["direct", "llm", "emb"] = "direct",
+        hint_index_path: str = None,
+        hint_retriever_path: str = None,
+        hint_num_results: int = 5,
+        skip_hints_for_current_task: bool = False,
+        hint_retrieval_mode: Literal["direct", "llm", "emb"] = "direct",
         hint_level: Literal["episode", "step"] = "episode",
         queries: list[str] | None = None,
     ) -> None:
         super().__init__(visible=use_task_hint)
         self.use_task_hint = use_task_hint
+        self.hint_type = hint_type
+        self.hint_index_type = hint_index_type
+        self.hint_query_type = hint_query_type
+        self.hint_index_path = hint_index_path
+        self.hint_retriever_path = hint_retriever_path
+        self.hint_num_results = hint_num_results
         self.hint_db_rel_path = "hint_db.csv"
         self.hint_db_path = hint_db_path  # Allow external path override
         self.hint_retrieval_mode: Literal["direct", "llm", "emb"] = hint_retrieval_mode
@@ -343,29 +369,47 @@ def __init__(
     def _init(self):
         """Initialize the block."""
         try:
-            # Use external path if provided, otherwise fall back to relative path
-            if self.hint_db_path and Path(self.hint_db_path).exists():
-                hint_db_path = Path(self.hint_db_path)
+            if self.hint_type == "docs":
+                if self.hint_index_type == "sparse":
+                    print("Loading sparse hint index")
+                    import bm25s
+                    self.hint_index = bm25s.BM25.load(self.hint_index_path, load_corpus=True)
+                    print("Sparse hint index loaded successfully")
+                elif self.hint_index_type == "dense":
+                    print("Loading dense hint index and retriever")
+                    from datasets import load_from_disk
+                    from sentence_transformers import SentenceTransformer
+                    self.hint_index = load_from_disk(self.hint_index_path)
+                    self.hint_index.load_faiss_index("embeddings", self.hint_index_path.removesuffix("/") + ".faiss")
+                    print("Dense hint index loaded successfully")
+                    self.hint_retriever = SentenceTransformer(self.hint_retriever_path)
+                    print("Hint retriever loaded successfully")
+                else:
+                    raise ValueError(f"Unknown hint index type: {self.hint_index_type}")
             else:
-                hint_db_path = Path(__file__).parent / self.hint_db_rel_path
-
-            if hint_db_path.exists():
-                self.hint_db = pd.read_csv(hint_db_path, header=0, index_col=None, dtype=str)
-                # Verify the expected columns exist
-                if "task_name" not in self.hint_db.columns or "hint" not in self.hint_db.columns:
-                    print(
-                        f"Warning: Hint database missing expected columns. Found: {list(self.hint_db.columns)}"
-                    )
+                # Use external path if provided, otherwise fall back to relative path
+                if self.hint_db_path and Path(self.hint_db_path).exists():
+                    hint_db_path = Path(self.hint_db_path)
+                else:
+                    hint_db_path = Path(__file__).parent / self.hint_db_rel_path
+
+                if hint_db_path.exists():
+                    self.hint_db = pd.read_csv(hint_db_path, header=0, index_col=None, dtype=str)
+                    # Verify the expected columns exist
+                    if "task_name" not in self.hint_db.columns or "hint" not in self.hint_db.columns:
+                        print(
+                            f"Warning: Hint database missing expected columns. Found: {list(self.hint_db.columns)}"
+                        )
+                        self.hint_db = pd.DataFrame(columns=["task_name", "hint"])
+                else:
+                    print(f"Warning: Hint database not found at {hint_db_path}")
                     self.hint_db = pd.DataFrame(columns=["task_name", "hint"])
-            else:
-                print(f"Warning: Hint database not found at {hint_db_path}")
-                self.hint_db = pd.DataFrame(columns=["task_name", "hint"])
-
-            self.hints_source = HintsSource(
-                hint_db_path=hint_db_path.as_posix(),
-                hint_retrieval_mode=self.hint_retrieval_mode,
-                skip_hints_for_current_task=self.skip_hints_for_current_task,
-            )
+                    
+                self.hints_source = HintsSource(
+                    hint_db_path=hint_db_path.as_posix(),
+                    hint_retrieval_mode=self.hint_retrieval_mode,
+                    skip_hints_for_current_task=self.skip_hints_for_current_task,
+                )
         except Exception as e:
             # Fallback to empty database on any error
             print(f"Warning: Could not load hint database: {e}")
@@ -376,6 +420,32 @@ def get_hints_for_task(self, task_name: str) -> str:
         if not self.use_task_hint:
             return ""
 
+        if self.hint_type == "docs":
+            if not hasattr(self, "hint_index"):
+                self._init()
+
+            if self.hint_query_type == "goal":
+                query = self.goal
+            elif self.hint_query_type == "llm":
+                query = self.llm.generate(self._prompt + self._abstract_ex + self._concrete_ex)
+            else:
+                raise ValueError(f"Unknown hint query type: {self.hint_query_type}")
+
+            if self.hint_index_type == "sparse":
+                query_tokens = bm25s.tokenize(query)
+                docs = self.hint_index.search(query_tokens, k=self.hint_num_results)
+                docs = docs["text"]
+            elif self.hint_index_type == "dense":
+                query_embedding = self.hint_retriever.encode(query)
+                _, docs = self.hint_index.get_nearest_examples("embeddings", query_embedding, k=self.hint_num_results)
+                docs = docs["text"]
+
+            hints_str = (
+                "# Hints:\nHere are some hints for the task you are working on:\n"
+                + "\n".join(docs)
+            )
+            return hints_str
+
         # Ensure hint_db is initialized
         if not hasattr(self, "hint_db"):
             self._init()

src/agentlab/agents/hitl_agent/base_multi_candidate_agent.py

@@ -0,0 +1,50 @@
+from typing_extensions import Protocol
+
+from agentlab.agents.agent_args import AgentArgs
+
+
+class MultiCandidateAgent(Protocol):
+    """
+    Protocol for agents that generate multiple candidates for get_action.
+
+    This protocol defines the contract for agents that can generate
+    multiple candidate actions and allow selection of one of them for execution.
+    """
+
+    def get_candidate_generations(
+        self, obs: dict, hint: list[str] | None = None, n_candidates: int = 3
+    ) -> "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
+        in a single step.
+
+        Args:
+            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
+        """
+        ...
+
+    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):
+    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:
+            self.agent_name = "MC-" + self.agent_name