app/retrieval/hybrid.py

"""Reciprocal rank fusion for vector + BM25 retrieval lists."""

from __future__ import annotations

from rank_bm25 import BM25Okapi

from app.models.schemas import SearchResult

_RRF_K = 60


def tokenize(text: str) -> list[str]:
    return [t for t in text.lower().split() if t]


def build_bm25_index(texts: list[str]) -> BM25Okapi | None:
    corpus = [tokenize(t) for t in texts]
    if not corpus:
        return None
    return BM25Okapi(corpus)


def bm25_search(
    query: str,
    *,
    texts: list[str],
    meta_rows: list[SearchResult],
    k: int,
) -> list[SearchResult]:
    """Return top-k BM25 hits from ``meta_rows`` aligned with ``texts``."""
    if not texts or not meta_rows:
        return []
    index = build_bm25_index(texts)
    if index is None:
        return []
    scores = index.get_scores(tokenize(query))
    ranked = sorted(
        zip(meta_rows, scores, strict=True),
        key=lambda pair: float(pair[1]),
        reverse=True,
    )
    out: list[SearchResult] = []
    for row, score in ranked[:k]:
        out.append(row.model_copy(update={"score": float(score)}))
    return out


def reciprocal_rank_fusion(
    ranked_lists: list[list[SearchResult]],
    *,
    top_n: int,
    rrf_k: int = _RRF_K,
) -> list[SearchResult]:
    """Merge multiple ranked lists with RRF; higher score = better."""
    scores: dict[str, float] = {}
    by_id: dict[str, SearchResult] = {}
    for lst in ranked_lists:
        for rank, item in enumerate(lst, start=1):
            scores[item.chunk_id] = scores.get(item.chunk_id, 0.0) + 1.0 / (rrf_k + rank)
            by_id[item.chunk_id] = item
    ordered = sorted(scores.keys(), key=lambda cid: scores[cid], reverse=True)
    return [
        by_id[cid].model_copy(update={"score": scores[cid]})
        for cid in ordered[:top_n]
    ]