154 lines
5.2 KiB
Python
154 lines
5.2 KiB
Python
from __future__ import annotations
|
|
|
|
from uuid import NAMESPACE_URL, uuid5
|
|
|
|
import pytest
|
|
|
|
from ragcore.retrieval.ranking import FusionConfig, fuse_candidates
|
|
from ragcore.retrieval.search.models import (
|
|
Candidate,
|
|
CandidateBackendKind,
|
|
CandidateProvenance,
|
|
)
|
|
|
|
|
|
def candidate(
|
|
backend: CandidateBackendKind,
|
|
label: str,
|
|
rank: int,
|
|
*,
|
|
document: str = "document-a",
|
|
parent: str | None = None,
|
|
) -> Candidate:
|
|
document_id = uuid5(NAMESPACE_URL, document)
|
|
version_id = uuid5(NAMESPACE_URL, f"{document}:version")
|
|
return Candidate(
|
|
point_id=uuid5(NAMESPACE_URL, f"point:{label}:{backend}"),
|
|
score=1.0,
|
|
rank=rank,
|
|
backend=backend,
|
|
provenance=CandidateProvenance(
|
|
workspace_id=uuid5(NAMESPACE_URL, "workspace"),
|
|
space_id=uuid5(NAMESPACE_URL, "space"),
|
|
document_id=document_id,
|
|
document_version_id=version_id,
|
|
chunk_id=uuid5(NAMESPACE_URL, f"chunk:{label}"),
|
|
source_id=uuid5(NAMESPACE_URL, f"source:{document}"),
|
|
parent_chunk_id=uuid5(NAMESPACE_URL, parent) if parent else None,
|
|
status="published",
|
|
content_type="text/plain",
|
|
language="en",
|
|
trust_level="source",
|
|
),
|
|
)
|
|
|
|
|
|
def test_rrf_matches_golden_example_and_keeps_explicit_score_semantics() -> None:
|
|
result = fuse_candidates(
|
|
[
|
|
candidate(CandidateBackendKind.DENSE, "a", 1),
|
|
candidate(CandidateBackendKind.DENSE, "b", 2),
|
|
],
|
|
[
|
|
candidate(CandidateBackendKind.SPARSE, "b", 1),
|
|
candidate(CandidateBackendKind.SPARSE, "c", 2),
|
|
],
|
|
)
|
|
|
|
assert [item.candidate.provenance.chunk_id for item in result] == [
|
|
uuid5(NAMESPACE_URL, "chunk:b"),
|
|
uuid5(NAMESPACE_URL, "chunk:a"),
|
|
uuid5(NAMESPACE_URL, "chunk:c"),
|
|
]
|
|
assert result[0].rrf_score == pytest.approx(1 / 61 + 1 / 62)
|
|
assert result[0].dense_rank == 2
|
|
assert result[0].sparse_rank == 1
|
|
assert not hasattr(result[0], "probability")
|
|
|
|
|
|
def test_ties_have_stable_identity_order_independent_of_input_order() -> None:
|
|
dense = [
|
|
candidate(CandidateBackendKind.DENSE, "z", 1),
|
|
candidate(CandidateBackendKind.DENSE, "a", 1),
|
|
]
|
|
first = fuse_candidates(dense, [])
|
|
second = fuse_candidates(list(reversed(dense)), [])
|
|
|
|
assert first == second
|
|
assert [item.candidate.provenance.chunk_id for item in first] == sorted(
|
|
(uuid5(NAMESPACE_URL, "chunk:z"), uuid5(NAMESPACE_URL, "chunk:a")),
|
|
key=str,
|
|
)
|
|
|
|
|
|
def test_duplicate_and_near_duplicate_chunks_are_suppressed() -> None:
|
|
duplicate_dense = candidate(CandidateBackendKind.DENSE, "duplicate", 1)
|
|
duplicate_sparse = candidate(CandidateBackendKind.SPARSE, "duplicate", 4)
|
|
near_first = candidate(CandidateBackendKind.DENSE, "near-first", 2, parent="parent-a")
|
|
near_second = candidate(CandidateBackendKind.DENSE, "near-second", 3, parent="parent-a")
|
|
|
|
result = fuse_candidates(
|
|
[duplicate_dense, near_first, near_second],
|
|
[duplicate_sparse],
|
|
)
|
|
|
|
assert len(result) == 2
|
|
assert result[0].dense_rank == 1
|
|
assert result[0].sparse_rank == 4
|
|
assert result[1].candidate.provenance.chunk_id == near_first.provenance.chunk_id
|
|
|
|
|
|
def test_per_document_limit_is_enforced_before_fused_limit() -> None:
|
|
result = fuse_candidates(
|
|
[
|
|
candidate(CandidateBackendKind.DENSE, "a1", 1, document="doc-a"),
|
|
candidate(CandidateBackendKind.DENSE, "a2", 2, document="doc-a"),
|
|
candidate(CandidateBackendKind.DENSE, "a3", 3, document="doc-a"),
|
|
candidate(CandidateBackendKind.DENSE, "b1", 4, document="doc-b"),
|
|
],
|
|
[],
|
|
config=FusionConfig(max_fused_candidates=4, max_per_document=2),
|
|
)
|
|
|
|
assert len(result) == 3
|
|
assert (
|
|
sum(
|
|
item.candidate.provenance.document_id == uuid5(NAMESPACE_URL, "doc-a")
|
|
for item in result
|
|
)
|
|
== 2
|
|
)
|
|
assert any(
|
|
item.candidate.provenance.document_id == uuid5(NAMESPACE_URL, "doc-b") for item in result
|
|
)
|
|
|
|
|
|
def test_config_bounds_reject_unbounded_or_invalid_values() -> None:
|
|
with pytest.raises(ValueError):
|
|
FusionConfig(rrf_k=0)
|
|
with pytest.raises(ValueError):
|
|
FusionConfig(max_per_document=5, max_fused_candidates=4)
|
|
with pytest.raises(ValueError):
|
|
FusionConfig(rrf_k=1.5) # type: ignore[arg-type]
|
|
with pytest.raises(ValueError):
|
|
FusionConfig(max_per_document=True) # type: ignore[arg-type]
|
|
|
|
|
|
def test_conflicting_provenance_and_backend_list_mismatch_fail_closed() -> None:
|
|
with pytest.raises(ValueError, match="provenance"):
|
|
fuse_candidates(
|
|
[candidate(CandidateBackendKind.DENSE, "same", 1)],
|
|
[candidate(CandidateBackendKind.SPARSE, "same", 1, document="doc-b")],
|
|
)
|
|
with pytest.raises(ValueError, match="backend"):
|
|
fuse_candidates([candidate(CandidateBackendKind.SPARSE, "wrong", 1)], [])
|
|
|
|
|
|
def test_input_candidate_bound_is_enforced() -> None:
|
|
candidates = [
|
|
candidate(CandidateBackendKind.DENSE, f"chunk-{index}", index + 1) for index in range(51)
|
|
]
|
|
|
|
with pytest.raises(ValueError, match="bounded"):
|
|
fuse_candidates(candidates, [])
|