Skip to content

Commit

Permalink
fix: improve unpacking of keyword search results
Browse files Browse the repository at this point in the history
  • Loading branch information
lsorber committed Nov 22, 2024
1 parent fdf803b commit e3d7537
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 2 deletions.
5 changes: 3 additions & 2 deletions src/raglite/_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]


Expand Down
11 changes: 11 additions & 0 deletions tests/test_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit e3d7537

Please sign in to comment.