Fixed pagination issue in search.py and added corresponding test (#846)

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

paperqa/agents/search.py

@@ -399,10 +399,12 @@ async def query(
         addresses = [
             s[1]
             for s in searcher.search(
-                index.parse_query(self.clean_query(query), query_fields), top_n
+                index.parse_query(self.clean_query(query), query_fields),
+                top_n,
+                offset=offset,
             ).hits
             if s[0] > min_score
-        ][offset : offset + top_n]
+        ]
         search_index_docs = [searcher.doc(address) for address in addresses]
         return [
             result

tests/test_agents.py

@@ -1061,6 +1061,25 @@ async def test_clinical_tool_usage(agent_test_settings) -> None:
     ), "No clinical trials were put into contexts"
 
 
+@pytest.mark.asyncio
+async def test_search_pagination(agent_test_settings: Settings) -> None:
+    """Test that pagination works correctly in SearchIndex.query()."""
+    index = await get_directory_index(settings=agent_test_settings)
+
+    page_size = 1
+
+    page1_results = await index.query(query="test", top_n=page_size, offset=0)
+    page2_results = await index.query(query="test", top_n=page_size, offset=page_size)
+    page1and2_results = await index.query(query="test", top_n=2 * page_size, offset=0)
+
+    assert (
+        page1_results == page1and2_results[:page_size]
+    ), "First page should match start of all results"
+    assert (
+        page2_results == page1and2_results[page_size : page_size * 2]
+    ), "Second page should match second slice of all results"
+
+
 class TestClinicalTrialSearchTool:
     @pytest.mark.asyncio
     async def test_continuation(self) -> None: