Caching the search index within a PaperSearch tool, to only call get_…

…directory_index per trajectory, with a test

paperqa/agents/tools.py

@@ -7,14 +7,14 @@
 import sys
 from typing import ClassVar, cast
 
-from pydantic import BaseModel, ConfigDict, Field, computed_field
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, computed_field
 
 from paperqa.docs import Docs
 from paperqa.llms import EmbeddingModel, LiteLLMModel
 from paperqa.settings import Settings
 from paperqa.types import Answer, DocDetails
 
-from .search import get_directory_index
+from .search import SearchIndex, get_directory_index
 
 logger = logging.getLogger(__name__)
 
@@ -82,6 +82,7 @@ class PaperSearch(NamedTool):
     settings: Settings
     embedding_model: EmbeddingModel
     previous_searches: dict[tuple[str, str | None], int] = Field(default_factory=dict)
+    _search_index: SearchIndex | None = PrivateAttr(default=None)
 
     async def paper_search(
         self,
@@ -125,12 +126,15 @@ async def paper_search(
             offset = self.previous_searches[search_key] = 0
 
         logger.info(f"Starting paper search for {query!r}.")
-        index = await get_directory_index(settings=self.settings, build=False)
-        results: list[Docs] = await index.query(
+        if self._search_index is None:
+            self._search_index = await get_directory_index(
+                settings=self.settings, build=False
+            )
+        results: list[Docs] = await self._search_index.query(
             query,
             top_n=self.settings.agent.search_count,
             offset=offset,
-            field_subset=[f for f in index.fields if f != "year"],
+            field_subset=[f for f in self._search_index.fields if f != "year"],
         )
         logger.info(
             f"{self.TOOL_FN_NAME} for query {query!r} and offset {offset} returned"

tests/test_agents.py

@@ -23,6 +23,7 @@
 from ldp.llms import EmbeddingModel, MultipleCompletionLLMModel
 from pydantic import ValidationError
 from pytest_subtests import SubTests
+from tantivy import Index
 
 from paperqa.agents import SearchIndex, agent_query
 from paperqa.agents.env import settings_to_tools
@@ -240,7 +241,11 @@ async def test_agent_types(
         " accept the answer for now, as we're in debug mode."
     )
     request = QueryRequest(query=question, settings=agent_test_settings)
-    response = await agent_query(request, agent_type=agent_type)
+    with patch.object(
+        Index, "open", side_effect=Index.open, autospec=True
+    ) as mock_open:
+        response = await agent_query(request, agent_type=agent_type)
+    mock_open.assert_called_once()
     assert response.answer.answer, "Answer not generated"
     assert response.answer.answer != "I cannot answer", "Answer not generated"
     assert response.answer.context, "No contexts were found"
@@ -462,20 +467,34 @@ async def test_agent_sharing_state(
         search_tool = PaperSearch(
             settings=agent_test_settings, embedding_model=embedding_model
         )
-        with patch.object(
-            SearchIndex, "save_index", autospec=True, wraps=SearchIndex.save_index
-        ) as mock_save_index:
+        with (
+            patch.object(
+                SearchIndex, "save_index", wraps=SearchIndex.save_index, autospec=True
+            ) as mock_save_index,
+            patch.object(
+                Index, "open", side_effect=Index.open, autospec=True
+            ) as mock_open,
+        ):
             await search_tool.paper_search(
                 "XAI self explanatory model",
                 min_year=None,
                 max_year=None,
                 state=env_state,
             )
-        assert env_state.docs.docs, "Search did not add any papers"
-        mock_save_index.assert_not_awaited(), "Search shouldn't try to update the index"
-        assert all(
-            isinstance(d, Doc) for d in env_state.docs.docs.values()
-        ), "Document type or DOI propagation failure"
+            assert env_state.docs.docs, "Search did not add any papers"
+            mock_open.assert_called_once()
+            assert all(
+                isinstance(d, Doc) for d in env_state.docs.docs.values()
+            ), "Document type or DOI propagation failure"
+
+            await search_tool.paper_search(
+                "XAI for chemical property prediction",
+                min_year=2018,
+                max_year=2024,
+                state=env_state,
+            )
+            mock_open.assert_called_once()
+            mock_save_index.assert_not_awaited()
 
     with subtests.test(msg=GatherEvidence.__name__):
         assert not answer.contexts, "No contexts is required for a later assertion"