DataOps-platform/app/core/knowledge/retrieval/fusion.py

from __future__ import annotations

from collections import defaultdict
from collections.abc import Mapping, Sequence

from app.core.knowledge.retrieval.contracts import KnowledgeEvidence


def reciprocal_rank_fusion(
    ranked: Mapping[str, Sequence[KnowledgeEvidence]],
    *,
    limit: int = 30,
    k: int = 60,
) -> tuple[KnowledgeEvidence, ...]:
    scores: dict[str, float] = defaultdict(float)
    evidence_by_id: dict[str, KnowledgeEvidence] = {}
    provenance: dict[str, list[str]] = defaultdict(list)
    for retriever, values in ranked.items():
        for rank, evidence in enumerate(values, start=1):
            scores[evidence.chunk_id] += 1.0 / (k + rank)
            evidence_by_id.setdefault(evidence.chunk_id, evidence)
            if retriever not in provenance[evidence.chunk_id]:
                provenance[evidence.chunk_id].append(retriever)
    ordered = sorted(scores, key=lambda key: (-scores[key], key))[:limit]
    return tuple(
        evidence_by_id[key].with_score_and_retriever(
            scores[key], "+".join(provenance[key])
        )
        for key in ordered
    )