58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import pytest
|
|
|
|
from rag.fusion import reciprocal_rank_fusion
|
|
from rag.models import RetrievalDocument, SearchHit, SourceRef
|
|
|
|
|
|
SOURCE = SourceRef(physical_page=1, precision="page")
|
|
|
|
|
|
def _hit(doc_id: str, score: float = 1.0) -> SearchHit:
|
|
return SearchHit(
|
|
RetrievalDocument(
|
|
doc_id=doc_id,
|
|
drug_id="drug",
|
|
kind="prose",
|
|
text=doc_id,
|
|
section_key="section",
|
|
source_refs=(SOURCE,),
|
|
),
|
|
score,
|
|
)
|
|
|
|
|
|
def test_rrf_promotes_candidate_supported_by_dense_and_lexical_rankings() -> None:
|
|
fused = reciprocal_rank_fusion(
|
|
[[_hit("dense-only"), _hit("shared")], [_hit("shared"), _hit("lexical-only")]],
|
|
rank_constant=60,
|
|
)
|
|
assert [hit.document.doc_id for hit in fused] == [
|
|
"shared", "dense-only", "lexical-only",
|
|
]
|
|
|
|
|
|
def test_rrf_deduplicates_a_document_within_one_ranking() -> None:
|
|
fused = reciprocal_rank_fusion(
|
|
[[_hit("duplicate"), _hit("duplicate")], [_hit("other")]],
|
|
rank_constant=10,
|
|
)
|
|
duplicate = next(hit for hit in fused if hit.document.doc_id == "duplicate")
|
|
assert duplicate.score == pytest.approx(1 / 11)
|
|
|
|
|
|
def test_rrf_limit_and_ties_are_deterministic() -> None:
|
|
fused = reciprocal_rank_fusion(
|
|
[[_hit("first")], [_hit("second")]], rank_constant=60, limit=1,
|
|
)
|
|
assert [hit.document.doc_id for hit in fused] == ["first"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("kwargs", "message"),
|
|
[({"rank_constant": 0}, "rank_constant"), ({"limit": -1}, "limit")],
|
|
)
|
|
def test_rrf_rejects_invalid_configuration(kwargs, message: str) -> None:
|
|
with pytest.raises(ValueError, match=message):
|
|
reciprocal_rank_fusion([], **kwargs)
|
|
|