56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
"""Rank fusion primitives for hybrid and multi-query retrieval.
|
|
|
|
This module is deliberately store-agnostic: dense, lexical, and rewritten-query
|
|
retrievers only need to return ranked ``SearchHit`` lists. Keeping fusion pure
|
|
makes its ordering deterministic and lets the live service add parallel I/O
|
|
without coupling domain code to Qdrant or PostgreSQL.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
|
|
from .models import SearchHit
|
|
|
|
|
|
def reciprocal_rank_fusion(
|
|
rankings: Sequence[Sequence[SearchHit]],
|
|
*,
|
|
rank_constant: int = 60,
|
|
limit: int | None = None,
|
|
) -> list[SearchHit]:
|
|
"""Fuse ranked candidate lists using reciprocal-rank fusion (RRF).
|
|
|
|
A document contributes at most once per input ranking. The returned score
|
|
is the accumulated RRF score, not a provider-specific similarity score, so
|
|
dense and lexical results remain comparable without score normalization.
|
|
Ties are stable by first appearance, which keeps results reproducible.
|
|
"""
|
|
if rank_constant < 1:
|
|
raise ValueError("rank_constant must be positive")
|
|
if limit is not None and limit < 0:
|
|
raise ValueError("limit must be non-negative or None")
|
|
|
|
scores: dict[str, float] = {}
|
|
documents = {}
|
|
first_seen: dict[str, int] = {}
|
|
seen_order = 0
|
|
|
|
for ranking in rankings:
|
|
seen_in_ranking: set[str] = set()
|
|
for rank, hit in enumerate(ranking, start=1):
|
|
doc_id = hit.document.doc_id
|
|
if doc_id in seen_in_ranking:
|
|
continue
|
|
seen_in_ranking.add(doc_id)
|
|
if doc_id not in documents:
|
|
documents[doc_id] = hit.document
|
|
first_seen[doc_id] = seen_order
|
|
seen_order += 1
|
|
scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (rank_constant + rank)
|
|
|
|
ordered_ids = sorted(scores, key=lambda doc_id: (-scores[doc_id], first_seen[doc_id]))
|
|
if limit is not None:
|
|
ordered_ids = ordered_ids[:limit]
|
|
return [SearchHit(document=documents[doc_id], score=scores[doc_id]) for doc_id in ordered_ids]
|
|
|