Skip to content

Commit 33a6412

Browse files
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>
1 parent 3ce7801 commit 33a6412

File tree

2 files changed

+23
-2
lines changed

2 files changed

+23
-2
lines changed

paperqa/agents/search.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -399,10 +399,12 @@ async def query(
399399
addresses = [
400400
s[1]
401401
for s in searcher.search(
402-
index.parse_query(self.clean_query(query), query_fields), top_n
402+
index.parse_query(self.clean_query(query), query_fields),
403+
top_n,
404+
offset=offset,
403405
).hits
404406
if s[0] > min_score
405-
][offset : offset + top_n]
407+
]
406408
search_index_docs = [searcher.doc(address) for address in addresses]
407409
return [
408410
result

tests/test_agents.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1061,6 +1061,25 @@ async def test_clinical_tool_usage(agent_test_settings) -> None:
10611061
), "No clinical trials were put into contexts"
10621062

10631063

1064+
@pytest.mark.asyncio
1065+
async def test_search_pagination(agent_test_settings: Settings) -> None:
1066+
"""Test that pagination works correctly in SearchIndex.query()."""
1067+
index = await get_directory_index(settings=agent_test_settings)
1068+
1069+
page_size = 1
1070+
1071+
page1_results = await index.query(query="test", top_n=page_size, offset=0)
1072+
page2_results = await index.query(query="test", top_n=page_size, offset=page_size)
1073+
page1and2_results = await index.query(query="test", top_n=2 * page_size, offset=0)
1074+
1075+
assert (
1076+
page1_results == page1and2_results[:page_size]
1077+
), "First page should match start of all results"
1078+
assert (
1079+
page2_results == page1and2_results[page_size : page_size * 2]
1080+
), "Second page should match second slice of all results"
1081+
1082+
10641083
class TestClinicalTrialSearchTool:
10651084
@pytest.mark.asyncio
10661085
async def test_continuation(self) -> None:

0 commit comments

Comments
 (0)