fix: align query tool name

Author: rejojer

openkb/agent/query.py

@@ -30,7 +30,7 @@
 """
 
 
-def pageindex_retrieve(doc_id: str, question: str, okb_dir: str, model: str) -> str:
+def _pageindex_retrieve_impl(doc_id: str, question: str, okb_dir: str, model: str) -> str:
     """Retrieve relevant content from a long document via PageIndex.
 
     Args:
@@ -148,7 +148,7 @@ def read_file(path: str) -> str:
         return read_wiki_file(path, wiki_root)
 
     @function_tool
-    def retrieve(doc_id: str, question: str) -> str:
+    def pageindex_retrieve(doc_id: str, question: str) -> str:
         """Retrieve relevant content from a long document via PageIndex.
 
         Use this when you need detailed content from a document that was
@@ -158,12 +158,12 @@ def retrieve(doc_id: str, question: str) -> str:
             doc_id: PageIndex document identifier (found in index.md).
             question: The question you are trying to answer.
         """
-        return pageindex_retrieve(doc_id, question, okb_dir, model)
+        return _pageindex_retrieve_impl(doc_id, question, okb_dir, model)
 
     return Agent(
         name="wiki-query",
         instructions=instructions,
-        tools=[list_files, read_file, retrieve],
+        tools=[list_files, read_file, pageindex_retrieve],
         model=model,
     )
 

tests/test_query.py

@@ -6,7 +6,7 @@
 
 import pytest
 
-from openkb.agent.query import build_query_agent, pageindex_retrieve, run_query
+from openkb.agent.query import _pageindex_retrieve_impl, build_query_agent, run_query
 from openkb.schema import SCHEMA_MD
 
 
@@ -24,7 +24,13 @@ def test_agent_tool_names(self, tmp_path):
         names = {t.name for t in agent.tools}
         assert "list_files" in names
         assert "read_file" in names
-        assert "retrieve" in names
+        assert "pageindex_retrieve" in names
+
+    def test_instructions_reference_registered_pageindex_tool(self, tmp_path):
+        agent = build_query_agent(str(tmp_path), str(tmp_path / "pi"), "gpt-4o-mini")
+        tool_names = {t.name for t in agent.tools}
+        assert "pageindex_retrieve" in agent.instructions
+        assert "pageindex_retrieve" in tool_names
 
     def test_schema_in_instructions(self, tmp_path):
         agent = build_query_agent(str(tmp_path), str(tmp_path / "pi"), "gpt-4o-mini")
@@ -63,7 +69,7 @@ def test_returns_page_content(self, tmp_path):
             mock_llm.return_value = MagicMock(
                 choices=[MagicMock(message=MagicMock(content="1-2"))]
             )
-            result = pageindex_retrieve("doc123", "What is the intro?", "/db", "gpt-4o-mini")
+            result = _pageindex_retrieve_impl("doc123", "What is the intro?", "/db", "gpt-4o-mini")
 
         assert "Introduction text here." in result
         assert "More intro content." in result
@@ -76,7 +82,7 @@ def test_handles_empty_structure(self, tmp_path):
         mock_client.collection.return_value = mock_col
 
         with patch("openkb.agent.query.PageIndexClient", return_value=mock_client):
-            result = pageindex_retrieve("doc456", "What?", "/db", "gpt-4o-mini")
+            result = _pageindex_retrieve_impl("doc456", "What?", "/db", "gpt-4o-mini")
 
         assert "No structure found" in result
 
@@ -88,7 +94,7 @@ def test_handles_structure_error(self, tmp_path):
         mock_client.collection.return_value = mock_col
 
         with patch("openkb.agent.query.PageIndexClient", return_value=mock_client):
-            result = pageindex_retrieve("doc789", "What?", "/db", "gpt-4o-mini")
+            result = _pageindex_retrieve_impl("doc789", "What?", "/db", "gpt-4o-mini")
 
         assert "Error" in result