src/generate.py

"""
Stage 6 — grounded generation with a local model (DeepSeek-R1-Distill-7B via Ollama).

WHY grounding is the whole point: a raw LLM answers from memory and will confidently
invent a wrong derivation. RAG constrains it to answer ONLY from retrieved chunks, cite
which chunk each claim came from, and refuse ("I don't know") when the material doesn't
cover the question. For a probability course that's the difference between a study aid and
a hallucination machine. We also gate on the reranker score: if the best retrieved chunk
is weak, we refuse rather than generate from irrelevant context.

DeepSeek-R1 emits its reasoning inside <think>...</think> then the answer. We split them:
the answer is the deliverable; the thinking is shown optionally (it's great for studying).
"""
from __future__ import annotations

import os
import re
import json
import urllib.request
import urllib.error

from src.store import search, search_multi  # imports torch/FlagEmbedding before qdrant (ordering matters)

_OLLAMA_URL = "http://localhost:11434/api/chat"

# Model registry: short name -> where/how to call it. "api" picks the wire format in _chat().
# WHY native Ollama /api/chat (not its OpenAI-compatible /v1): /v1 ignores num_ctx, and
# num_ctx=4096 is exactly what keeps the retrieval models + LLM co-resident on the 8GB GPU.
MODELS: dict[str, dict] = {
    # qwen stays the default so LOCAL behavior is byte-identical when no new env vars are set.
    "qwen": {"model": "qwen2.5:3b-instruct", "url": _OLLAMA_URL, "api": "ollama",
             "api_key_env": None},
    "glm-local": {"model": "glm4:9b", "url": _OLLAMA_URL, "api": "ollama",
                  "api_key_env": None},
    "glm-flash": {"model": os.environ.get("ZAI_MODEL", "glm-4.7-flash"),
                  "url": "https://api.z.ai/api/paas/v4/chat/completions", "api": "openai",
                  "api_key_env": "ZAI_API_KEY"},
}

# PROBRAG_MODEL is a legacy override (README documents PROBRAG_MODEL=deepseek-r1:7b): it swaps
# the model id behind the "qwen" name rather than adding an entry, so old invocations still work.
_legacy_model = os.environ.get("PROBRAG_MODEL")
if _legacy_model:
    MODELS["qwen"]["model"] = _legacy_model

# PROBRAG_MODELS = comma-separated names to ENABLE (default just qwen -> local unchanged).
ENABLED_MODELS = [
    n.strip() for n in os.environ.get("PROBRAG_MODELS", "qwen").split(",")
    if n.strip() in MODELS
] or ["qwen"]
DEFAULT_MODEL = os.environ.get("PROBRAG_DEFAULT_MODEL") or ENABLED_MODELS[0]
if DEFAULT_MODEL not in ENABLED_MODELS:
    DEFAULT_MODEL = ENABLED_MODELS[0]


def _resolve(model: str | None) -> dict:
    """Map a requested model name to its registry entry. None -> default. Invalid or
    not-enabled -> ValueError (server maps to 400)."""
    name = model or DEFAULT_MODEL
    if name not in ENABLED_MODELS:
        raise ValueError(f"unknown model {name!r}; enabled: {ENABLED_MODELS}")
    return MODELS[name]
MIN_RERANK_SCORE = 0.005           # low backstop only. The eval-derived 0.024 cleanly split the
                                   # GOLDEN phrasings, but real-world phrasings of answerable
                                   # questions can score lower and get wrongly refused. So we keep
                                   # the hard threshold minimal (catch ~zero-signal retrieval) and
                                   # let the GROUNDED PROMPT be the real guard — it refuses when the
                                   # retrieved context doesn't actually contain the answer.
# On 8GB VRAM the embed+rerank models (~2.2GB) and the 7B (~4.7GB) must coexist on the GPU.
# A smaller context shrinks the 7B's KV cache so everything fits on-GPU (fast generation)
# instead of spilling Ollama to CPU. 4K is plenty for top_k grounded passages.
NUM_CTX = 4096
CTX_TOKEN_BUDGET = 2600            # room for system + question + answer within NUM_CTX
MAX_PARENT_TOKENS = 1100           # truncate oversized parents (chunker max ~7600)
_CHARS_PER_TOK = 4

SYSTEM = (
    "You are a probability tutor. Answer the question using ONLY the numbered context "
    "passages — no outside facts. Be concise and plain: the shortest answer that is still "
    "correct, everyday words, no filler and no unnecessary jargon. Keep the essential terms "
    "the course tests (variance, moment generating function, Bernoulli trials) but gloss each "
    "in plain words the first time it appears. Textbook results often appear inside worked "
    "examples and in notation; treat these as answers (g(t) = moment generating function; "
    "V(X) or sigma^2 = variance; E(X) or mu = expected value; a 'first success' count is "
    "geometric; a Bernoulli-trials success count is binomial). A worked Example or derivation "
    "IS a valid answer: if a passage shows the steps (e.g. computes g(t), then g'(0) and "
    "g''(0), to get a mean or variance), walk through those steps and state the result. "
    "Extract the key formula and cite its tag, e.g. [S1]. Keep LaTeX as written ($...$, "
    "$$...$$). Only if the passages are unrelated to the question, reply exactly: \"I don't "
    "know based on the provided material.\""
)


# Only comparison / derive-from questions need splitting; everything else takes the
# single-query path untouched (keeps the 39 saturated single-hop cases byte-identical, and
# avoids paying an LLM call + N× retrieval on queries one embedding already covers).
DECOMPOSE_GATE = re.compile(
    r"\bcompare\b|\bversus\b|\bvs\b|difference between|\bderive\b.*\b(from|via|using|through)\b", re.I)

DECOMPOSE_SYSTEM = (
    "Split a probability question into focused search queries, ONE per line, no numbering "
    "and no other text. Rules: (1) each query targets a SINGLE quantity of a SINGLE "
    "distribution so it can be searched on its own; (2) always name the distribution IN FULL "
    "— write 'the geometric distribution', never bare 'geometric' (which collides with the "
    "geometric mean); (3) if the question asks about two quantities (e.g. mean and variance) "
    "of two distributions, emit one query per (quantity, distribution) pair. Example — for "
    "'compare the mean and variance of the Poisson and geometric distributions':\n"
    "mean of the Poisson distribution\n"
    "variance of the Poisson distribution\n"
    "mean of the geometric distribution\n"
    "variance of the geometric distribution\n"
    "Output only the queries."
)


def decompose(question: str) -> list[str]:
    """Split a multi-part question into focused sub-queries via the local LLM, but ONLY when
    the gate matches (comparison / derive-from). Returns [] otherwise, or if the model
    returns nothing usable — the caller then falls back to plain single-query retrieval.
    Temp 0 so the split is stable enough for the eval harness to A/B against."""
    if not DECOMPOSE_GATE.search(question):
        return []
    try:
        # Prefer the local qwen entry so sub-query splitting doesn't burn hosted-API quota.
        # When qwen isn't enabled (the deployed Space has no Ollama), use the default hosted
        # entry instead — a decompose call is ~100 tokens, and losing decomposition would
        # silently degrade every comparison question (the 45/45 multi-hop wins).
        entry = MODELS["qwen"] if "qwen" in ENABLED_MODELS else _resolve(None)
        raw = _chat([{"role": "system", "content": DECOMPOSE_SYSTEM},
                     {"role": "user", "content": question}], entry)
    except Exception:
        return []  # fail open: a decompose hiccup falls back to single-query retrieval
    subs = []
    for line in raw.splitlines():
        s = re.sub(r"^\s*(?:[-*\d.)]+\s*)+", "", line).strip()  # strip bullets/numbering
        if s:
            subs.append(s)
    return subs[:4]  # up to 4: a "mean and variance of A and B" question is 4 (quantity, dist) pairs


def _retrieve_for(client, question: str, top_k: int) -> list[dict]:
    """Retrieve grounding hits: multi-query when the question decomposes, else single-query
    rerank. The one entry point so answer/stream/why all get decomposition for free."""
    subs = decompose(question)
    if subs:
        # Sub-query lanes only — the original comparison query as a lane just re-adds the
        # single-topic-biased hits decomposition exists to get past, and costs a rerank pass.
        return search_multi(client, subs, top_k=top_k)
    return search(client, question, top_k=top_k, mode="rerank")


def _truncate(text: str, max_tokens: int) -> str:
    cap = max_tokens * _CHARS_PER_TOK
    if len(text) <= cap:
        return text
    return text[:cap].rsplit("\n", 1)[0] + "\n... [passage truncated]"


def _build_context(hits: list[dict]) -> tuple[str, list[dict]]:
    """Pack passages under a token budget so the prompt fits NUM_CTX. Oversized parents
    are truncated; once the budget is hit we stop adding passages."""
    blocks, sources = [], []
    used = 0
    for i, h in enumerate(hits, 1):
        text = _truncate(h["parent_text"], MAX_PARENT_TOKENS)
        toks = len(text) // _CHARS_PER_TOK
        if used + toks > CTX_TOKEN_BUDGET and blocks:
            break  # budget exhausted; keep what we have
        used += toks
        tag = f"S{i}"
        loc = " / ".join(x for x in (h.get("chapter"), h.get("section")) if x) or h["source"]
        blocks.append(f"[{tag}] ({loc})\n{text}")
        sources.append({"tag": tag, "chunk_id": h["chunk_id"], "location": loc,
                        "page": h.get("page"),  # 1-indexed PDF page -> UI opens the viewer here
                        "score": h.get("rerank_score", h.get("score"))})
    return "\n\n".join(blocks), sources


def _headers(entry: dict) -> dict:
    h = {"Content-Type": "application/json"}
    key_env = entry.get("api_key_env")
    if key_env:
        key = os.environ.get(key_env)
        if not key:
            raise RuntimeError(f"model {entry['model']!r} needs env {key_env}, which is unset")
        h["Authorization"] = f"Bearer {key}"
    return h


def _post(entry: dict, payload: dict):
    req = urllib.request.Request(
        entry["url"], data=json.dumps(payload).encode(), headers=_headers(entry))
    try:
        return urllib.request.urlopen(req, timeout=600)
    except urllib.error.HTTPError as e:
        # Surface the provider's error body (e.g. Z.ai "model not found") — without this the
        # SSE stream just truncates silently and the failure is undiagnosable in production.
        body = e.read()[:300].decode("utf-8", "replace")
        raise RuntimeError(f"LLM API error {e.code} from {entry['model']}: {body}") from e


def _chat(messages: list[dict], entry: dict) -> str:
    """Non-streaming completion, dispatching on entry['api']."""
    if entry["api"] == "openai":
        payload = {"model": entry["model"], "messages": messages, "stream": False,
                   "temperature": 0.0}
        with _post(entry, payload) as r:
            return json.loads(r.read())["choices"][0]["message"]["content"]
    # ollama: native NDJSON /api/chat. num_ctx keeps the models co-resident on the 8GB GPU.
    payload = {"model": entry["model"], "messages": messages, "stream": False,
               "options": {"num_ctx": NUM_CTX, "temperature": 0.0}}
    with _post(entry, payload) as r:
        return json.loads(r.read())["message"]["content"]


def _chat_stream(messages: list[dict], entry: dict):
    """Yield content deltas, dispatching on entry['api']."""
    if entry["api"] == "openai":
        payload = {"model": entry["model"], "messages": messages, "stream": True,
                   "temperature": 0.0}
        with _post(entry, payload) as r:
            for raw in r:
                line = raw.decode("utf-8").strip()
                if not line.startswith("data:"):
                    continue
                data = line[len("data:"):].strip()
                if data == "[DONE]":
                    return
                # choices can be empty on trailing usage/keep-alive chunks — don't crash the stream
                choices = json.loads(data).get("choices") or [{}]
                delta = choices[0].get("delta", {}).get("content", "")
                if delta:
                    yield delta
        return
    # ollama: NDJSON stream.
    payload = {"model": entry["model"], "messages": messages, "stream": True,
               "options": {"num_ctx": NUM_CTX, "temperature": 0.0}}
    with _post(entry, payload) as r:
        for line in r:
            chunk = json.loads(line)
            piece = chunk.get("message", {}).get("content", "")
            if piece:
                yield piece
            if chunk.get("done"):
                return


def _split_think(text: str) -> tuple[str, str]:
    think = "\n".join(re.findall(r"<think>(.*?)</think>", text, flags=re.DOTALL)).strip()
    answer = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
    return answer, think


WHY_SYSTEM = (
    "You are a probability tutor giving a DEEPER explanation of a previous answer. Use ONLY "
    "the numbered context passages — no outside facts. The student already has the short "
    "answer and now wants the reasoning: explain the intuition, walk through the derivation "
    "step by step, define the symbols, and note where each step comes from. Use plain words "
    "and no unnecessary jargon — gloss each technical term in plain language as you use it. "
    "Same notation conventions as before (g(t) = mgf, V(X) = variance, etc.). Cite passages "
    "as [S1]. Keep LaTeX as written ($...$, $$...$$). If the context doesn't support a deeper "
    "explanation, say so plainly rather than inventing one."
)

ELI5_SYSTEM = (
    "You are a probability tutor explaining to a total beginner. Use ONLY the numbered "
    "context passages — no outside facts. Explain in the FEWEST words possible: simplest "
    "everyday language, short sentences, no jargon and no filler. Use one everyday analogy "
    "only if it truly helps. Keep the actual formula (in $...$) so the answer stays correct, "
    "but say what it means in plain words right after. Cite the passage tag, e.g. [S1]. If "
    "the context doesn't cover it, say so plainly."
)

# style -> (system prompt, user-instruction prefix). Keeps stream() one code path.
_STYLES = {
    "default": (SYSTEM, "Question:"),
    "deeper": (WHY_SYSTEM, "Explain in depth, with reasoning:"),
    "eli5": (ELI5_SYSTEM, "Explain as simply as possible, fewest words:"),
}


def stream(client, question: str, *, top_k: int = 6, show_thinking: bool = False,
           style: str = "default", model: str | None = None):
    """Streaming twin of answer()/why(). Yields event dicts: {"type": "sources", ...} first
    (so the UI can render citations before generation), then {"type": "token", "text": ...}
    deltas, then {"type": "done", ...} carrying the cleaned final answer — the UI replaces
    the streamed text with it.
    # ponytail: under PROBRAG_MODEL=deepseek-r1 the <think> block streams visibly and is
    # only stripped in the final event; filter deltas live if that ever matters (default
    # qwen2.5:3b never emits think tags)."""
    entry = _resolve(model)  # raises ValueError before retrieval on a bad name
    hits = _retrieve_for(client, question, top_k)
    best = hits[0].get("rerank_score", 0.0) if hits else 0.0
    if not hits or best < MIN_RERANK_SCORE:
        yield {"type": "done", "answer": "I don't know based on the provided material.",
               "sources": [], "refused": True, "best_score": best}
        return
    context, sources = _build_context(hits)
    yield {"type": "sources", "sources": sources, "best_score": best}
    system, prefix = _STYLES.get(style, _STYLES["default"])
    user = f"Context passages:\n\n{context}\n\n{prefix} {question}"
    parts: list[str] = []
    for piece in _chat_stream([{"role": "system", "content": system},
                               {"role": "user", "content": user}], entry):
        parts.append(piece)
        yield {"type": "token", "text": piece}
    ans, think = _split_think("".join(parts))
    done = {"type": "done", "answer": ans, "refused": False}
    if show_thinking:
        done["thinking"] = think
    yield done


def answer(client, question: str, *, top_k: int = 6, show_thinking: bool = False,
           model: str | None = None) -> dict:
    entry = _resolve(model)  # raises ValueError before retrieval on a bad name
    hits = _retrieve_for(client, question, top_k)
    best = hits[0].get("rerank_score", 0.0) if hits else 0.0
    if not hits or best < MIN_RERANK_SCORE:
        return {"answer": "I don't know based on the provided material.",
                "sources": [], "refused": True, "best_score": best}

    context, sources = _build_context(hits)
    user = f"Context passages:\n\n{context}\n\nQuestion: {question}"
    raw = _chat([{"role": "system", "content": SYSTEM},
                 {"role": "user", "content": user}], entry)
    ans, think = _split_think(raw)
    out = {"answer": ans, "sources": sources, "refused": False, "best_score": best}
    if show_thinking:
        out["thinking"] = think
    return out