File tree Expand file tree Collapse file tree 2 files changed +23
-2
lines changed Expand file tree Collapse file tree 2 files changed +23
-2
lines changed Original file line number Diff line number Diff line change @@ -399,10 +399,12 @@ async def query(
399
399
addresses = [
400
400
s [1 ]
401
401
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 ,
403
405
).hits
404
406
if s [0 ] > min_score
405
- ][ offset : offset + top_n ]
407
+ ]
406
408
search_index_docs = [searcher .doc (address ) for address in addresses ]
407
409
return [
408
410
result
Original file line number Diff line number Diff line change @@ -1061,6 +1061,25 @@ async def test_clinical_tool_usage(agent_test_settings) -> None:
1061
1061
), "No clinical trials were put into contexts"
1062
1062
1063
1063
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
+
1064
1083
class TestClinicalTrialSearchTool :
1065
1084
@pytest .mark .asyncio
1066
1085
async def test_continuation (self ) -> None :
You can’t perform that action at this time.
0 commit comments