pdf-quiz-generator/backend/app/services/rerank_service.py

"""Cross-encoder reranking: the model reads the query and the document together.

Everything upstream of here ranks without ever putting the two side by side. A
bi-encoder embedded each row long before anybody typed anything, so cosine
distance measures "same topic", which is not the same as "answers this". BM25
measures which words happen to coincide. Reciprocal rank fusion combines those
two opinions but cannot add a third.

A cross-encoder is that third opinion: one forward pass over the pair, so it can
tell that a section titled "Croup" which spends its whole body on epiglottitis
is not the thing to put first. The price is that it cannot be precomputed — the
work is per pair, per query — which is why nothing here ever sees the corpus.
It sees a shortlist the cheap rankers already produced, and reorders it.

**Reranking is an opinion about order and nothing else.** Every caller keeps the
rows retrieval found; a reranker that is unconfigured, down, slow or talking
nonsense returns `None` from `rerank()` and the caller ships today's ordering.
A search page with no results because a scoring service is unwell is a much
worse failure than a search page in a slightly worse order, and the second
failure is invisible to a user while the first ends their session.

Served through the same proxy as everything else (`LITELLM_API_BASE`), so a
learner's query goes where the embeddings already go and nowhere new. Nothing
here knows which reranker is behind that name: `POST /v1/rerank` with
`{model, query, documents, top_n}`, back `{results: [{index, relevance_score}]}`,
is the shape Cohere defined and Jina, BGE and LiteLLM's own rerank route all
speak. This deployment's proxy serves `cohere-rerank-v4.0-pro` and
`-fast` alongside `jina-reranker-v2-base-multilingual`, which is open-weight —
moving to a self-hosted one is a change to `LITELLM_RERANK_MODEL` and no change
to any code.
"""
import hashlib
import json
import logging

from app.config import settings

logger = logging.getLogger(__name__)

#: Per-document budget. Rerankers truncate at roughly 1k tokens per document
#: anyway, and the part of a clinical vignette that decides relevance is its
#: opening, not the option list; sending more buys nothing and costs latency
#: linear in total tokens.
DOC_CHARS = 1200

#: A cross-encoder is trained on query-length queries. Handing it a whole
#: uploaded document (which `search_service._query_terms` exists to cope with)
#: is out of distribution as well as slow, so the query is clamped too.
QUERY_CHARS = 2000

#: Measured against this proxy: 10 candidates ≈ 250 ms, 30 ≈ 380 ms, 50 ≈ 430 ms,
#: 100 ≈ 580 ms — mostly the round trip, then roughly 2 ms a pair. The fused pool
#: is 60 wide and a page shows ten, so past about fifty the reranker is being
#: paid to confirm that rows neither cheap ranker wanted are still not wanted.
#: The numbers and the method are in docs/reranking.md.
MAX_CANDIDATES = 50

#: Deterministic in (model, query, documents), and the document text is part of
#: the key, so an edited row invalidates its own entry. That leaves nothing a
#: day can make wrong — the same reasoning as the query-embedding cache.
CACHE_TTL = 24 * 3600

#: Long enough for a cold model to answer, short enough that a wedged proxy
#: costs one page load rather than the request. On timeout the caller keeps the
#: order it already had.
TIMEOUT_SECONDS = 8


def rerank_model() -> str:
    """The active reranker, Redis override first — same shape as the embedder.

    An administrator can move between served rerankers without a redeploy;
    unset everywhere means the feature is simply off.
    """
    try:
        import redis as redis_lib

        client = redis_lib.from_url(settings.REDIS_URL, decode_responses=True,
                                    socket_connect_timeout=1)
        chosen = client.get("settings:rerank_model")
        if chosen is not None:
            return chosen.strip()
    except Exception:
        pass
    return (settings.LITELLM_RERANK_MODEL or "").strip()


def is_configured() -> bool:
    """Whether reranking can be attempted at all."""
    return bool(rerank_model() and settings.LITELLM_API_KEY and settings.LITELLM_API_BASE)


def _cache_key(model: str, query: str, documents: list[str]) -> str:
    digest = hashlib.sha256()
    digest.update(query.encode("utf-8", "ignore"))
    for document in documents:
        digest.update(b"\x00")
        digest.update(document.encode("utf-8", "ignore"))
    return f"rerank:{model}:{digest.hexdigest()[:40]}"


def _redis():
    try:
        import redis as redis_lib

        return redis_lib.from_url(settings.REDIS_URL, decode_responses=True,
                                  socket_connect_timeout=1)
    except Exception:
        return None


def _parse(payload, count: int) -> list[float] | None:
    """Scores aligned to the documents that went in, or None if they are not.

    Positional trust is the bug waiting to happen here: the response is sorted
    by score, so reading it in order attaches the best document's score to the
    first document. Every score is placed by the `index` the server echoes, and
    a response that does not cover every input is refused outright rather than
    half-applied — a partial reordering is indistinguishable from a bad one.
    """
    if not isinstance(payload, dict):
        return None
    results = payload.get("results")
    if not isinstance(results, list) or len(results) != count:
        return None
    scores: list[float | None] = [None] * count
    for item in results:
        if not isinstance(item, dict):
            return None
        index, score = item.get("index"), item.get("relevance_score")
        if not isinstance(index, int) or not isinstance(score, (int, float)):
            return None
        if not 0 <= index < count or scores[index] is not None:
            return None
        scores[index] = float(score)
    if any(score is None for score in scores):
        return None
    return scores  # type: ignore[return-value]


def rerank(query: str, documents: list[str]) -> list[float] | None:
    """Relevance of each document to the query, in the order they were given.

    `None` means "no opinion" and is the answer to every kind of trouble: not
    configured, unreachable, slow, or a reply that does not line up with the
    request. Callers must read it as "keep what you had".
    """
    query = " ".join((query or "").split())[:QUERY_CHARS]
    if not query or not documents:
        return None
    model = rerank_model()
    if not is_configured():
        return None

    clipped = [" ".join((document or "").split())[:DOC_CHARS] or " " for document in documents]
    key = _cache_key(model, query, clipped)
    cache = _redis()
    if cache is not None:
        try:
            hit = cache.get(key)
            if hit:
                cached = json.loads(hit)
                if isinstance(cached, list) and len(cached) == len(clipped):
                    return [float(score) for score in cached]
        except Exception:
            logger.debug("Rerank cache unreadable", exc_info=True)

    api_base = (settings.LITELLM_API_BASE or "").rstrip("/").removesuffix("/v1")
    try:
        import httpx

        response = httpx.post(
            f"{api_base}/v1/rerank",
            headers={"Authorization": f"Bearer {settings.LITELLM_API_KEY}",
                     "Content-Type": "application/json"},
            # `top_n` is the whole list on purpose. Asking for fewer would save
            # nothing — the pairs are scored either way — and would leave the
            # tail unscored, so it could only be appended in its old order.
            json={"model": model, "query": query, "documents": clipped,
                  "top_n": len(clipped)},
            timeout=TIMEOUT_SECONDS,
        )
        response.raise_for_status()
        scores = _parse(response.json(), len(clipped))
    except Exception as error:
        logger.warning("Rerank unavailable (%s); keeping fused order", error)
        return None

    if scores is None:
        logger.warning("Rerank returned a reply that did not match the request; "
                       "keeping fused order")
        return None

    if cache is not None:
        try:
            cache.setex(key, CACHE_TTL, json.dumps(scores))
        except Exception:
            logger.debug("Could not cache rerank scores", exc_info=True)
    return scores


def reorder(query: str, items: list, text_of) -> list:
    """`items` best-first by cross-encoder score, unchanged if there is no score.

    Only the head is reordered. Beyond `MAX_CANDIDATES` the list is passed
    through untouched, which keeps the cost of a query independent of how many
    rows matched it and keeps pagination past the first page stable.

    The return value is always a permutation of the input: nothing is dropped
    and nothing is added, whatever the reranker does or fails to do. That is the
    property that makes this safe to put in front of a search page.
    """
    if len(items) < 2:
        return list(items)
    head = list(items[:MAX_CANDIDATES])
    tail = list(items[len(head):])
    scores = rerank(query, [text_of(item) for item in head])
    # The length is `rerank`'s contract, and a caller that stubs or replaces it
    # is exactly when a reordering half the list long would go unnoticed.
    if scores is None or len(scores) != len(head):
        return list(items)
    # Ties keep the fused order, so agreement between the cheap rankers still
    # settles anything the cross-encoder is indifferent about.
    order = sorted(range(len(head)), key=lambda index: (-scores[index], index))
    return [head[index] for index in order] + tail