(wip) refactor hinting index

Author: patricebechard

src/agentlab/agents/generic_agent_hinter/generic_agent.py

@@ -16,12 +16,14 @@
 import bgym
 from bgym import Benchmark
 from browsergym.experiments.agent import Agent, AgentInfo
-
+import pandas as pd
+from pathlib import Path
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.agents.agent_args import AgentArgs
 from agentlab.llm.chat_api import BaseModelArgs
 from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
 from agentlab.llm.tracking import cost_tracker_decorator
+from agentlab.agents.tool_use_agent.tool_use_agent import HintsSource
 
 from .generic_agent_prompt import (
     GenericPromptFlags,
@@ -92,6 +94,8 @@ def __init__(
         self.action_set = self.flags.action.action_set.make_action_set()
         self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)
 
+        self._init_hints_index()
+
         self._check_flag_constancy()
         self.reset(seed=None)
 
@@ -246,3 +250,46 @@ def _get_maxes(self):
             else 20  # dangerous to change the default value here?
         )
         return max_prompt_tokens, max_trunc_itr
+
+    def _init_hints_index(self):
+        """Initialize the block."""
+        try:
+            if self.flags.hint_type == "docs":
+                if self.flags.hint_index_type == "sparse":
+                    import bm25s
+                    self.hint_index = bm25s.BM25.load(self.flags.hint_index_path, load_corpus=True)
+                elif self.flags.hint_index_type == "dense":
+                    from datasets import load_from_disk
+                    from sentence_transformers import SentenceTransformer
+                    self.hint_index = load_from_disk(self.flags.hint_index_path)
+                    self.hint_index.load_faiss_index("embeddings", self.flags.hint_index_path.removesuffix("/") + ".faiss")
+                    self.hint_retriever = SentenceTransformer(self.flags.hint_retriever_path)
+                else:
+                    raise ValueError(f"Unknown hint index type: {self.flags.hint_index_type}")
+            else:
+                # Use external path if provided, otherwise fall back to relative path
+                if self.flags.hint_db_path and Path(self.flags.hint_db_path).exists():
+                    hint_db_path = Path(self.flags.hint_db_path)
+                else:
+                    hint_db_path = Path(__file__).parent / self.flags.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"])
+                self.hints_source = HintsSource(
+                    hint_db_path=hint_db_path.as_posix(),
+                    hint_retrieval_mode=self.flags.hint_retrieval_mode,
+                    skip_hints_for_current_task=self.flags.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}")
+            self.hint_db = pd.DataFrame(columns=["task_name", "hint"])

src/agentlab/agents/generic_agent_hinter/generic_agent_prompt.py

@@ -80,6 +80,7 @@ def __init__(
         actions: list[str],
         memories: list[str],
         thoughts: list[str],
+        hints: list[str],
         previous_plan: str,
         step: int,
         flags: GenericPromptFlags,
@@ -120,6 +121,7 @@ def time_for_caution():
         self.think = dp.Think(visible=lambda: flags.use_thinking)
         self.hints = dp.Hints(visible=lambda: flags.use_hints)
         goal_str: str = goal[0]["text"]
+        # TODO: This design is not very good as we will instantiate the loop up at every step
         self.task_hint = TaskHint(
             use_task_hint=flags.use_task_hint,
             hint_db_path=flags.hint_db_path,
@@ -147,7 +149,8 @@ def _prompt(self) -> HumanMessage:
 
         # Add task hints if enabled
         task_hints_text = ""
-        if self.flags.use_task_hint and hasattr(self, "task_name"):
+        # if self.flags.use_task_hint and hasattr(self, "task_name"):
+        if self.flags.use_task_hint:
             task_hints_text = self.task_hint.get_hints_for_task(self.task_name)
 
         prompt.add_text(
@@ -371,19 +374,14 @@ def _init(self):
         try:
             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:
@@ -422,8 +420,8 @@ def get_hints_for_task(self, task_name: str) -> str:
 
         if self.hint_type == "docs":
             if not hasattr(self, "hint_index"):
+                print("Initializing hint index new time")
                 self._init()
-
             if self.hint_query_type == "goal":
                 query = self.goal
             elif self.hint_query_type == "llm":
@@ -432,9 +430,15 @@ def get_hints_for_task(self, task_name: str) -> str:
                 raise ValueError(f"Unknown hint query type: {self.hint_query_type}")
 
             if self.hint_index_type == "sparse":
+                import bm25s
                 query_tokens = bm25s.tokenize(query)
-                docs = self.hint_index.search(query_tokens, k=self.hint_num_results)
-                docs = docs["text"]
+                docs, _ = self.hint_index.retrieve(query_tokens, k=self.hint_num_results)
+                docs = [elem["text"] for elem in docs[0]]
+                # HACK: truncate to 20k characters (should cover >99% of the cases)
+                for doc in docs:
+                    if len(doc) > 20000:
+                        doc = doc[:20000]
+                        doc += " ...[truncated]"
             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)