align query tool name

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