From e3d7537a8dc20ae30a17793e93ea14cc47daa5b0 Mon Sep 17 00:00:00 2001 From: Laurent Sorber Date: Fri, 22 Nov 2024 10:32:09 +0100 Subject: [PATCH] fix: improve unpacking of keyword search results --- src/raglite/_search.py | 5 +++-- tests/test_search.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/src/raglite/_search.py b/src/raglite/_search.py index c4f006a..374c3fd 100644 --- a/src/raglite/_search.py +++ b/src/raglite/_search.py @@ -129,8 +129,9 @@ def keyword_search( """) results = session.execute(statement, params={"match": fts5_query, "limit": num_results}) # Unpack the results. - chunk_ids, keyword_score = zip(*results, strict=True) - chunk_ids, keyword_score = list(chunk_ids), list(keyword_score) # type: ignore[assignment] + results = list(results) + chunk_ids = [result.chunk_id for result in results] + keyword_score = [result.score for result in results] return chunk_ids, keyword_score # type: ignore[return-value] diff --git a/tests/test_search.py b/tests/test_search.py index 8a626c0..9cea9d1 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -46,3 +46,14 @@ def test_search(raglite_test_config: RAGLiteConfig, search_method: SearchMethod) # Extend the chunks with their neighbours and group them into contiguous segments. segments = retrieve_segments(chunk_ids, neighbors=(-1, 1), config=raglite_test_config) assert all(isinstance(segment, str) for segment in segments) + + +def test_search_no_results(raglite_test_config: RAGLiteConfig, search_method: SearchMethod) -> None: + """Test searching for a query with no keyword search results.""" + query = "supercalifragilisticexpialidocious" + num_results = 5 + chunk_ids, scores = search_method(query, num_results=num_results, config=raglite_test_config) + num_results_expected = 0 if search_method == keyword_search else num_results + assert len(chunk_ids) == len(scores) == num_results_expected + assert all(isinstance(chunk_id, str) for chunk_id in chunk_ids) + assert all(isinstance(score, float) for score in scores)