src/retrieve/potion.py

"""Step 05 — potion broad retrieval: the recall pass with a static embedder, over the bm25_prefilter's output.

Embeds the pool's ``text`` ONLINE with the configured static embedder (``embedder_broad``, default
potion-base-32M — Model2Vec, a static token→vector lookup, ~13K docs/s, no transformer forward), scores each
candidate by cosine against the JD's ``ideal_summary`` (the subjective embedding query — HyDE was dropped),
and keeps the top ``top_n``. This is what lets the expensive granite embedder later touch only a few hundred
rows instead of the whole pool.

Deterministic: a static lookup with no RNG, plus a stable score+id sort — the surviving ids and
``potion_score`` are byte-reproducible across runs.
"""

from __future__ import annotations

import os
import time
from functools import lru_cache
from typing import Any

import numpy as np
import polars as pl

from common.logging import get_logger, step
from config import Settings, load_settings
from models.llm import MODEL_DIR  # shared in-image model cache dir (HF_HOME redirect)


@lru_cache(maxsize=1)
def _load_potion(repo: str) -> Any:
    """Load (downloading on first use) the static Model2Vec embedder; cached for the process."""
    os.environ.setdefault("HF_HOME", str(MODEL_DIR))  # resolve weights inside the repo, offline at rank time
    from model2vec import StaticModel  # imported lazily — keeps the import graph light

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    log = get_logger("retrieve")
    log.info("potion.load", repo=repo)
    start = time.perf_counter()
    # force_download=True is model2vec's DEFAULT — it skips the local-cache check entirely and always calls
    # snapshot_download, i.e. network on every load even when fully cached. False restores cache-first,
    # download-only-if-missing (see model2vec.persistence._resolve_folder).
    try:
        model = StaticModel.from_pretrained(repo, force_download=False)
    except Exception as exc:
        log.error("potion.load_failed", repo=repo, seconds=round(time.perf_counter() - start, 3), error=str(exc))
        raise
    log.info("potion.load_complete", repo=repo, seconds=round(time.perf_counter() - start, 3))
    return model


def _embed_unit(model: Any, texts: list[str]) -> np.ndarray:
    """Encode texts → L2-normalized float matrix (zero rows clamp to norm 1 so cosine is 0, never NaN)."""
    vectors = np.asarray(model.encode(texts), dtype=np.float32)
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return vectors / norms


def _narrow_by_cosine(candidates: pl.DataFrame, scores: np.ndarray, top_n: int) -> pl.DataFrame:
    """Attach ``potion_score`` and keep the top ``top_n`` (stable: score desc, then candidate_id)."""
    keep_n = min(max(1, top_n), len(scores))
    scored = candidates.with_columns(pl.Series("potion_score", scores.astype(float)))
    return scored.sort(["potion_score", "candidate_id"], descending=[True, False]).head(keep_n)


def potion_broad_retrieval(
    candidates: pl.DataFrame, query_text: str, settings: Settings | None = None
) -> pl.DataFrame:
    """Narrow the pool to the top-``top_n`` by potion cosine vs ``query_text`` (the JD ideal_summary).

    Adds a ``potion_score`` column. Disabled / empty query / empty pool ⇒ pass the full pool through with
    ``potion_score = 0.0`` (downstream stages then no-op on the flat column).
    """
    active = settings or load_settings()
    broad = active.broad_retrieval
    with step("rank.potion_broad", rows_in=candidates.height) as metrics:
        if not broad.enabled or not query_text.strip() or candidates.height == 0:
            metrics["rows_out"] = candidates.height
            metrics["skipped"] = True
            return candidates.with_columns(pl.lit(0.0).alias("potion_score"))
        model = _load_potion(active.models.embedder_broad)
        pool_vectors = _embed_unit(model, candidates.get_column("text").to_list())
        query_vector = _embed_unit(model, [query_text])[0]
        scores = pool_vectors @ query_vector  # cosine (both L2-normalized)
        narrowed = _narrow_by_cosine(candidates, scores, broad.top_n)
        metrics["rows_out"] = narrowed.height
        metrics["score_max"] = round(float(scores.max()), 4)
        return narrowed