src/store.py

"""
Stage 4 — embed (BGE-M3) + store (Qdrant), and Stage 5 — hybrid retrieval + rerank.

WHY BGE-M3: it emits a DENSE vector (semantic meaning: "variance of a sum" ~ "Var(X+Y)")
AND a SPARSE lexical vector (exact token/symbol overlap: matches "$\\sigma^2$", "binomial")
from a single model. Dense alone misses exact symbol matches; sparse alone misses
paraphrase. Fusing both (hybrid) is why this beats plain cosine on technical queries.

WHY parent/child here: we embed small retrieval units (children, or childless parents)
for precise matching, but store each unit's PARENT text in the payload so generation
gets the whole problem+solution, not a fragment.
"""
from __future__ import annotations

import os

# torch and FlagEmbedding both ship an OpenMP runtime (libiomp5md.dll); on Windows the
# duplicate triggers a native crash (0xC0000005). This makes the load robust.
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")

# IMPORT ORDER MATTERS: torch + FlagEmbedding native libs must load BEFORE qdrant_client,
# or the second one to load segfaults (0xC0000005) on Windows. Proven by isolation test.
import torch  # noqa: F401
import FlagEmbedding  # noqa: F401

import json
from pathlib import Path

from qdrant_client import QdrantClient, models

DENSE_DIM = 1024
COLLECTION = "probability"
_model = None
_reranker = None


def get_model():
    """Lazy-load BGE-M3. fp16 only when a GPU is visible; on CPU use fp32.
    The web server hides the GPU (CUDA_VISIBLE_DEVICES="") so the full 8GB is free for
    Ollama's 7B generation — retrieval on CPU is light (~1-3s)."""
    global _model
    if _model is None:
        cuda = torch.cuda.device_count() > 0
        # pass an explicit single device on CPU — auto-detect divides by GPU count (0 -> crash)
        _model = FlagEmbedding.BGEM3FlagModel(
            "BAAI/bge-m3", use_fp16=cuda, devices=None if cuda else "cpu")
    return _model


def get_reranker():
    """Lazy-load the bge-reranker-v2-m3 cross-encoder (fp16 on GPU, fp32 on CPU)."""
    global _reranker
    if _reranker is None:
        cuda = torch.cuda.device_count() > 0
        _reranker = FlagEmbedding.FlagReranker(
            "BAAI/bge-reranker-v2-m3", use_fp16=cuda, devices=None if cuda else "cpu")
    return _reranker


def embed(texts: list[str]) -> tuple[list[list[float]], list[dict]]:
    out = get_model().encode(
        texts, return_dense=True, return_sparse=True, return_colbert_vecs=False,
        batch_size=8, max_length=8192,
    )
    dense = [v.tolist() for v in out["dense_vecs"]]
    sparse = [{int(k): float(v) for k, v in lw.items()} for lw in out["lexical_weights"]]
    return dense, sparse


def _sparse_vec(d: dict[int, float]) -> models.SparseVector:
    return models.SparseVector(indices=list(d.keys()), values=list(d.values()))


def contextualize(chapter: str, section: str, text: str) -> str:
    """Prepend a one-line location header ("Chapter 10: Generating Functions > Examples")
    to a chunk before it is embedded and reranked. WHY: worked examples state results in
    notation, not words -- e.g. Example 10.4 computes the Poisson mgf as "g(t) = ..." but
    never writes the phrase "moment generating function", so a query using that phrase
    out-ranks it onto longer prose-theory chunks. The chapter title ("Generating
    Functions") supplies the missing vocabulary, letting the example surface for the query
    that it actually answers. This is lightweight contextual retrieval -- no LLM, computed
    once at index time, and it also makes the generation context self-labelling."""
    loc = " > ".join(p for p in (chapter, section) if p)
    return f"[{loc}]\n{text}" if loc else text


def retrieval_units(chunks: list[dict]) -> list[dict]:
    """Units to embed: children, plus parents that have no children. Each unit is
    annotated with the parent text it should return.

    Front/back matter (title page, Contents, Preface, Index, appendix data tables) carries
    an empty chapter label and is dropped here. WHY: the Index is one 7600-token chunk that
    lists every term in the book, so it out-ranks real content on keyword queries (it
    surfaced as the top hit for "Bayes' theorem", burying the actual Bayes' Formula section
    and causing a refusal). None of this matter is explanatory content, so excluding it
    strictly improves retrieval precision."""
    by_id = {c["id"]: c for c in chunks}
    parents_with_kids = {c["parent_id"] for c in chunks if c.get("parent_id")}
    units = []
    for c in chunks:
        if not c["chapter"]:
            continue  # front/back matter, not course content
        if c["is_parent"] and c["id"] in parents_with_kids:
            continue  # represented by its children
        parent = by_id.get(c.get("parent_id")) if c.get("parent_id") else c
        u = dict(c)
        u["parent_text"] = parent["text"]
        u["parent_id"] = parent["id"]
        units.append(u)
    return units


def build_collection(client: QdrantClient, name: str = COLLECTION) -> None:
    # Delete unconditionally rather than guarding on collection_exists(): points use
    # sequential int ids (0..N-1), so if a rebuild has FEWER units than before and the old
    # collection survives, upsert overwrites 0..N-1 and leaves the tail ids as stale points
    # (this once left the dropped book Index lingering and answerable). delete_collection is
    # a no-op tolerant call here.
    client.delete_collection(name)
    client.create_collection(
        name,
        vectors_config={"dense": models.VectorParams(size=DENSE_DIM, distance=models.Distance.COSINE)},
        sparse_vectors_config={"sparse": models.SparseVectorParams(modifier=models.Modifier.IDF)},
    )


def upsert_units(client: QdrantClient, units: list[dict], name: str = COLLECTION,
                 pages: dict[str, int] | None = None) -> int:
    # Embed the contextualized child text (location header + chunk). The reranker scores
    # against the contextualized PARENT text (stored below), so both retrieval stages see
    # the chapter/section vocabulary the worked examples lack.
    pages = pages or {}
    embed_texts = [contextualize(u["chapter"], u["section"], u["text"]) for u in units]
    dense, sparse = embed(embed_texts)
    points = []
    for i, u in enumerate(units):
        parent_text = contextualize(u["chapter"], u["section"], u["parent_text"])
        points.append(
            models.PointStruct(
                id=i,
                vector={"dense": dense[i], "sparse": _sparse_vec(sparse[i])},
                payload={
                    "chunk_id": u["id"],
                    "parent_id": u["parent_id"],
                    "type": u["type"],
                    "chapter": u["chapter"],
                    "section": u["section"],
                    "source": u["source"],
                    "text": u["text"],
                    "parent_text": parent_text,
                    "page": pages.get(u["parent_id"]),  # 1-indexed PDF page for the UI viewer
                },
            )
        )
    client.upsert(name, points)
    return len(points)


def _dedupe_parents(hits: list[dict]) -> list[dict]:
    """Keep the best-scoring unit per parent so the LLM sees distinct problems."""
    seen: dict[str, dict] = {}
    for h in hits:
        pid = h["parent_id"]
        if pid not in seen:
            seen[pid] = h
    return list(seen.values())


def rerank(query: str, hits: list[dict], *, top_k: int) -> list[dict]:
    """Cross-encoder rerank: score (query, parent_text) pairs, return top_k."""
    if not hits:
        return []
    # batch_size=8: FlagEmbedding's compute_score runs a throwaway "adjust batch size"
    # forward pass over the FIRST batch before the real loop — with the default (128) that
    # first batch is the whole pool, so every rerank did 2x the GPU work. A small batch
    # caps the waste at 8 pairs. Scores are identical (measured: 51 pairs 2.0s -> 1.06s).
    scores = get_reranker().compute_score(
        [[query, h["parent_text"]] for h in hits], normalize=True, batch_size=8)
    if not isinstance(scores, list):
        scores = [scores]
    for h, s in zip(hits, scores):
        h["rerank_score"] = float(s)
    hits.sort(key=lambda h: h["rerank_score"], reverse=True)
    return hits[:top_k]


def _hybrid_candidates(client: QdrantClient, query: str, cand: int, name: str) -> list[dict]:
    """One dense+sparse RRF retrieval → raw candidate payloads (pre-dedupe, pre-rerank)."""
    dense, sparse = embed([query])
    res = client.query_points(
        name,
        prefetch=[
            models.Prefetch(query=dense[0], using="dense", limit=cand),
            models.Prefetch(query=_sparse_vec(sparse[0]), using="sparse", limit=cand),
        ],
        query=models.FusionQuery(fusion=models.Fusion.RRF),
        limit=cand,
        with_payload=True,
    )
    return [{"score": p.score, **p.payload} for p in res.points]


def search(client: QdrantClient, query: str, *, top_k: int = 5, mode: str = "rerank",
           name: str = COLLECTION) -> list[dict]:
    """mode: 'dense' (vector-only), 'hybrid' (dense+sparse RRF), or 'rerank'
    (hybrid candidates re-scored by the cross-encoder)."""
    # Candidate pool feeds the cross-encoder, whose cost is linear (~40ms/pair on the
    # 4060, shared with Ollama) — the pool size IS the time-to-first-token. 80 candidates
    # cost ~5s/query; A/B on the 45-question golden set showed 24 is metric-identical
    # (Hit@1/3/5, MRR 0.881, same misses), so the old "big pool recovers vocabulary
    # mismatch" worry didn't survive measurement.
    cand = max(top_k * 4, 24)
    if mode == "dense":
        dense, _ = embed([query])
        res = client.query_points(name, query=dense[0], using="dense",
                                  limit=cand, with_payload=True)
        hits = _dedupe_parents([{"score": p.score, **p.payload} for p in res.points])
    else:
        hits = _dedupe_parents(_hybrid_candidates(client, query, cand, name))
    if mode == "rerank":
        return rerank(query, hits, top_k=top_k)
    return hits[:top_k]


def search_multi(client: QdrantClient, queries: list[str], *,
                 top_k: int = 5, cand: int = 40, name: str = COLLECTION) -> list[dict]:
    """Multi-query retrieval for comparison / multi-part questions. WHY: one query embedding
    can't sit near two distributions at once, so a comparison retrieves only one side.

    Each sub-query is retrieved AND reranked against ITSELF (a chunk about only the geometric
    distribution must be scored on "geometric variance", not on "compare binomial and
    geometric" — the latter buries single-topic chunks, which is exactly what killed the
    union-then-rerank approach). The per-query rankings are then round-robin merged, so every
    sub-query is guaranteed representation in the final top_k and both sides surface. LLM-free
    by design — the caller (generate.decompose) makes the sub-queries."""
    lanes = []
    for q in queries:
        hits = _dedupe_parents(_hybrid_candidates(client, q, cand, name))
        lanes.append(rerank(q, hits, top_k=top_k))  # scored against its OWN query
    merged, seen = [], set()
    for i in range(top_k):
        for lane in lanes:
            if i < len(lane) and lane[i]["parent_id"] not in seen:
                seen.add(lane[i]["parent_id"])
                merged.append(lane[i])
                if len(merged) >= top_k:
                    return merged
    return merged


def load_chunks(path: Path) -> list[dict]:
    return json.loads(Path(path).read_text(encoding="utf-8"))