storage/reranker.py

"""storage/reranker.py — cross-encoder reranking for RAG chunks.

Wraps BAAI/bge-reranker-base via sentence-transformers.CrossEncoder.
Loaded lazily on first call so app startup stays fast (~100 MB model
downloaded once into ~/.cache/huggingface).

Used by storage/vector_store.search to re-rank Chroma's top-k candidates.
"""
from __future__ import annotations

from typing import Optional

_MODEL_NAME = "BAAI/bge-reranker-base"
_model = None  # type: ignore[var-annotated]


def _get_model():
    """Lazy-load the cross-encoder. Returns None on import or load failure."""
    global _model
    if _model is None:
        try:
            from sentence_transformers import CrossEncoder
            _model = CrossEncoder(_MODEL_NAME)
        except Exception as exc:  # pragma: no cover — defensive, not testable
            import sys
            print(f"[reranker] failed to load {_MODEL_NAME}: {exc}", file=sys.stderr)
            _model = False  # sentinel: tried and failed
    return _model if _model is not False else None


def rerank(query: str, candidates: list[dict], top_k: int = 3) -> list[dict]:
    """Re-score (query, candidate.text) pairs with the cross-encoder, return top_k.

    Each candidate is a dict with at least a 'text' field (matches the format
    returned by storage.vector_store.search).

    Falls back to the original ordering (truncated to top_k) if:
      - candidates already <= top_k (no work to do)
      - the cross-encoder model failed to load
      - the predict call raises
    """
    if not candidates or len(candidates) <= top_k:
        return candidates[:top_k]

    model = _get_model()
    if model is None:
        return candidates[:top_k]

    try:
        pairs = [(query, c.get("text", "")) for c in candidates]
        scores = model.predict(pairs)
        ranked = sorted(zip(candidates, scores), key=lambda x: float(x[1]), reverse=True)
        return [c for c, _ in ranked[:top_k]]
    except Exception:
        return candidates[:top_k]


def warmup() -> None:
    """Eagerly load the cross-encoder model; call at app startup to avoid cold-start penalty."""
    _get_model()


def reset_for_test() -> None:
    """Reset the cached model — used by unit tests that monkeypatch the loader."""
    global _model
    _model = None