src/core/hybrid_retriever.py

"""Hybrid BM25 + dense retrieval with reciprocal rank fusion (RRF)."""

from __future__ import annotations

import hashlib
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

from src.core.bm25_search import BM25Search
from src.core.retrieval_result import RetrievalResult
from src.core.vector_search import VectorSearch


@dataclass
class FusionConfig:
    k_rrf: int = 60
    candidate_k_bm25: int = 50
    candidate_k_vector: int = 50
    parallel: bool = True
    cache_max_entries: int = 128


class _LRUCache:
    def __init__(self, max_entries: int) -> None:
        self._max = max(0, max_entries)
        self._store: OrderedDict[str, List[RetrievalResult]] = OrderedDict()

    def get(self, key: str) -> Optional[List[RetrievalResult]]:
        if self._max == 0:
            return None
        if key not in self._store:
            return None
        self._store.move_to_end(key)
        # Return shallow copies so callers cannot mutate cache entries
        return [RetrievalResult(**r.__dict__) for r in self._store[key]]

    def set(self, key: str, value: List[RetrievalResult]) -> None:
        if self._max == 0:
            return
        self._store[key] = [RetrievalResult(**r.__dict__) for r in value]
        self._store.move_to_end(key)
        while len(self._store) > self._max:
            self._store.popitem(last=False)


def reciprocal_rank_fusion(
    ranked_lists: List[List[str]],
    k_rrf: int = 60,
    weights: Optional[List[float]] = None,
) -> List[Tuple[str, float]]:
    """
    Standard RRF: score(d) = sum_i 1 / (k_rrf + rank_i(d)).
    Missing from a list means that list does not contribute for that document.
    """
    scores: Dict[str, float] = {}
    for idx, ids in enumerate(ranked_lists):
        weight = 1.0
        if weights and idx < len(weights):
            weight = weights[idx]
        for rank, doc_id in enumerate(ids, start=1):
            scores[doc_id] = scores.get(doc_id, 0.0) + weight * (1.0 / (k_rrf + rank))
    return sorted(scores.items(), key=lambda x: (-x[1], x[0]))


class HybridRetriever:
    """Combines BM25Search and VectorSearch; ranks with RRF."""

    def __init__(
        self,
        bm25_search: BM25Search,
        vector_search: VectorSearch,
        *,
        fusion_weights: Optional[Dict[str, float]] = None,
        fusion_config: Optional[FusionConfig] = None,
        enable_cache: bool = True,
    ) -> None:
        self.bm25 = bm25_search
        self.vector = vector_search
        self.fusion_weights = fusion_weights or {"vector": 0.6, "bm25": 0.4}
        self.config = fusion_config or FusionConfig()
        self._cache = _LRUCache(self.config.cache_max_entries) if enable_cache else _LRUCache(0)

    def reciprocal_rank_fusion(
        self,
        ranked_lists: List[List[str]],
        k_rrf: Optional[int] = None,
    ) -> List[Tuple[str, float]]:
        k = k_rrf if k_rrf is not None else self.config.k_rrf
        weights = [self.fusion_weights.get("bm25", 1.0), self.fusion_weights.get("vector", 1.0)]
        return reciprocal_rank_fusion(ranked_lists, k_rrf=k, weights=weights)

    @staticmethod
    def _cache_key_parts(
        bm25_query: str,
        vector_query: str,
        k: int,
        k_rrf: int,
        ck_bm25: int,
        ck_vec: int,
        collection: str,
    ) -> str:
        raw = f"{collection}|{k}|{k_rrf}|{ck_bm25}|{ck_vec}|BM25:{bm25_query}|VEC:{vector_query}"
        return hashlib.sha256(raw.encode("utf-8")).hexdigest()

    def _run_bm25(self, bm25_query: str, k: int) -> List[Dict[str, Any]]:
        return self.bm25.search(bm25_query, k=k)

    def _run_vector(self, vector_query: str, k: int, filters: Optional[Dict]) -> List[Dict[str, Any]]:
        return self.vector.search(vector_query, k=k, filters=filters)

    def retrieve(
        self,
        bm25_query: str,
        vector_query: str,
        k: int = 20,
        *,
        filters: Optional[Dict] = None,
        collection_name_for_cache: str = "documents",
    ) -> List[RetrievalResult]:
        cfg = self.config
        ck_b = max(k, cfg.candidate_k_bm25)
        ck_v = max(k, cfg.candidate_k_vector)

        cache_key = self._cache_key_parts(
            bm25_query, vector_query, k, cfg.k_rrf, ck_b, ck_v, collection_name_for_cache
        )
        cached = self._cache.get(cache_key)
        if cached is not None:
            return cached[:k]

        bm25_hits: List[Dict[str, Any]] = []
        vec_hits: List[Dict[str, Any]] = []

        if cfg.parallel:
            with ThreadPoolExecutor(max_workers=2) as ex:
                f_b = ex.submit(self._run_bm25, bm25_query, ck_b)
                f_v = ex.submit(self._run_vector, vector_query, ck_v, filters)
                bm25_hits = f_b.result()
                vec_hits = f_v.result()
        else:
            bm25_hits = self._run_bm25(bm25_query, ck_b)
            vec_hits = self._run_vector(vector_query, ck_v, filters)

        bm25_ids = [str(h["id"]) for h in bm25_hits]
        vec_ids = [str(h["id"]) for h in vec_hits]

        weights = [self.fusion_weights.get("bm25", 1.0), self.fusion_weights.get("vector", 1.0)]
        fused = reciprocal_rank_fusion([bm25_ids, vec_ids], k_rrf=cfg.k_rrf, weights=weights)

        by_id: Dict[str, Dict[str, Any]] = {}
        for h in bm25_hits:
            by_id[str(h["id"])] = {**h, "_from_bm25": True}
        for h in vec_hits:
            hid = str(h["id"])
            if hid not in by_id:
                by_id[hid] = {**h, "_from_bm25": False}
            else:
                # merge vector-specific fields
                entry = by_id[hid]
                if "distance" in h and "distance" not in entry:
                    entry["distance"] = h["distance"]
                entry["_from_vector"] = True

        bm25_rank: Dict[str, int] = {doc_id: i + 1 for i, doc_id in enumerate(bm25_ids)}
        vec_rank: Dict[str, int] = {doc_id: i + 1 for i, doc_id in enumerate(vec_ids)}

        results: List[RetrievalResult] = []
        for doc_id, fusion_score in fused[:k]:
            row = by_id.get(doc_id, {})
            text = row.get("text") or ""
            meta = row.get("metadata") or {}
            bm25_sc = row.get("score")
            dist = row.get("distance")
            vec_sim = (1.0 - float(dist)) if dist is not None else None

            sources: List[str] = []
            if doc_id in bm25_rank:
                sources.append("bm25")
            if doc_id in vec_rank:
                sources.append("vector")

            br = bm25_rank.get(doc_id)
            vr = vec_rank.get(doc_id)
            confidence = self._confidence(fusion_score, br, vr)

            results.append(
                RetrievalResult(
                    id=doc_id,
                    text=text,
                    metadata=meta if isinstance(meta, dict) else {},
                    fusion_score=fusion_score,
                    bm25_rank=br,
                    vector_rank=vr,
                    bm25_score=float(bm25_sc) if bm25_sc is not None else None,
                    vector_similarity=vec_sim,
                    sources=sources,
                    confidence=confidence,
                )
            )

        self._cache.set(cache_key, results)
        return [RetrievalResult(**r.__dict__) for r in results]

    @staticmethod
    def _confidence(fusion_score: float, bm25_rank: Optional[int], vec_rank: Optional[int]) -> float:
        """Heuristic 0..1 from RRF score and how strong each leg is."""
        parts: List[float] = []
        if bm25_rank is not None:
            parts.append(1.0 / bm25_rank)
        if vec_rank is not None:
            parts.append(1.0 / vec_rank)
        leg = sum(parts) / max(len(parts), 1) if parts else 0.0
        return min(1.0, 0.5 * fusion_score + 0.5 * leg)