webapp/local_llm.py

"""In-process llama.cpp LLM backend for running the model inside a ZeroGPU Space.

On Hugging Face ZeroGPU the GPU is attached only for the duration of a
``@spaces.GPU`` call, so the entire map-reduce for a summary or report must run
inside one such call: the GGUF is loaded onto the GPU at the top of the call,
used for every chunk, then the GPU is released when the call returns. We never
run a persistent server.

Off-Space there is no ZeroGPU runtime; ``spaces.GPU`` then degrades to a plain
pass-through decorator and the model runs on CPU (``LLAMA_N_GPU_LAYERS=0``),
which is exactly how the local CPU smoke test exercises this seam.

The chroma map-reduce is untouched: we just hand
:func:`chroma.summarize.summarize_text` and :func:`chroma.report.generate_report`
a ``complete`` callable backed by a llama.cpp model instead of the OpenAI client.

Two GGUFs are downloaded at startup and the user picks between them per request
(see MODELS below): "e4b" (Gemma 4 E4B, public, fast — the default) and "26b"
(Gemma 4 26B-A4B, private, higher quality).

Configuration (env):
    GGUF_E4B_REPO / GGUF_E4B_FILE   override the E4B model's repo/file.
    GGUF_26B_REPO / GGUF_26B_FILE   override the 26B model's repo/file.
    DEFAULT_MODEL           which model key to use when none is specified ("e4b").
    LLAMA_N_CTX             context window (default 8192).
    LLAMA_N_GPU_LAYERS      -1 = offload all layers to GPU (ZeroGPU); 0 = CPU.
    LLAMA_MAX_TOKENS        default max output tokens per call.
    GPU_DURATION_SUMMARIZE / GPU_DURATION_REPORT / GPU_DURATION_ANSWER
                            ZeroGPU window seconds per task (default 90 / 180 / 90).
    REPORT_MAX_SECTIONS     cap on chunks a report mines (default 30).
    HF_TOKEN                speeds up / authenticates model downloads; required
                            for the private 26B GGUF repo.
"""

from __future__ import annotations

import ctypes
import glob
import os
import queue
import site
import threading
from functools import lru_cache
from typing import Iterator

import spaces

from chroma import generate_report, summarize_text
from chroma.answer import (
    ANSWER_PROMPT,
    ANSWER_SYSTEM,
    assemble_context,
    budget_for,
    resolve_engine,
)


def _preload_cuda_libs() -> None:
    """Preload CUDA runtime libs (cudart/cublas) from the torch-provided nvidia
    pip packages so llama.cpp's prebuilt CUDA wheel resolves them regardless of
    the loader path. Harmless off-Space (no nvidia packages -> nothing loaded).
    Order matters: cudart, then cublasLt, then cublas."""
    roots = list(site.getsitepackages()) + [site.getusersitepackages()]
    for pat in ("libcudart.so*", "libnvrtc.so*", "libcublasLt.so*", "libcublas.so*"):
        for root in roots:
            for so in glob.glob(os.path.join(root, "nvidia", "**", "lib", pat), recursive=True):
                try:
                    ctypes.CDLL(so, mode=ctypes.RTLD_GLOBAL)
                except OSError:
                    pass


_preload_cuda_libs()

# Two interchangeable Gemma-4 GGUFs the user can pick between at request time:
#   "e4b" — Gemma 4 E4B (~5.3 GB, public ggml-org build). Fast to load + infer; the
#           default. Fits every GPU window comfortably.
#   "26b" — Gemma 4 26B-A4B (25.2B-param MoE, ~3.8B active; ~16.8 GB). Our own GGUF
#           conversion (see model/gemma4_gguf.py); the repo is PRIVATE, so the Space
#           needs HF_TOKEN set for hf_hub_download to authenticate. Higher quality,
#           slower (reloaded onto the GPU per call since ZeroGPU releases it between
#           calls). Each entry's repo/file is env-overridable.
def _model_env(key: str, repo: str, file: str, label: str) -> dict:
    up = key.upper()
    return {
        "repo": os.getenv(f"GGUF_{up}_REPO", repo),
        "file": os.getenv(f"GGUF_{up}_FILE", file),
        "label": label,
        # Per-model context override (0 = use the global LLAMA_N_CTX). The big 26B GGUF
        # leaves little VRAM for the KV cache on a 24 GB A10G, so it can cap its own
        # context (GGUF_26B_N_CTX) below the larger window the small E4B can afford.
        "n_ctx": int(os.getenv(f"GGUF_{up}_N_CTX", "0")),
    }


# Three in-process GGUFs the user picks between (all run locally on the GPU):
#   "e4b"  — Gemma 4 E4B (small, fast; Q8_0 is near-lossless and tiny).
#   "26b"  — Gemma 4 26B-A4B at Q4_K_M (the standard quant — good quality, fits easily).
#   "full" — Gemma 4 26B-A4B at Q8_0 (near-full quality; biggest/slowest). Each entry's
#            repo/file/n_ctx is env-overridable (GGUF_<KEY>_REPO/FILE/N_CTX).
# Our fine-tuned agent models (Gemma 4 fine-tuned for the ReAct tool-calling protocol;
# `high` adds a GRPO stage). All three are PUBLIC, so no HF_TOKEN is needed to fetch them.
# Each entry's repo/file/n_ctx stays env-overridable (GGUF_<KEY>_REPO/FILE/N_CTX).
MODELS: dict[str, dict] = {
    "e4b": _model_env("e4b", "build-small-hackathon/agenda-parser-lite",
                      "agenda-parser-lite-Q8_0.gguf", "agenda-parser-lite"),
    "26b": _model_env("26b", "build-small-hackathon/agenda-parser-medium",
                      "agenda-parser-medium-Q4_K_M.gguf", "agenda-parser-medium"),
    "full": _model_env("full", "build-small-hackathon/agenda-parser-high",
                       "agenda-parser-high-Q8_0.gguf", "agenda-parser-high"),
}
DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "e4b").strip().lower()
if DEFAULT_MODEL not in MODELS:
    DEFAULT_MODEL = "e4b"

N_CTX = int(os.getenv("LLAMA_N_CTX", "8192"))
N_GPU_LAYERS = int(os.getenv("LLAMA_N_GPU_LAYERS", "-1"))  # -1 = all on GPU
MAX_OUTPUT_TOKENS = int(os.getenv("LLAMA_MAX_TOKENS", "1024"))


def resolve_model(model: str | None) -> str:
    """Normalize a request's model choice to a known key (falls back to default)."""
    key = (model or "").strip().lower()
    return key if key in MODELS else DEFAULT_MODEL


def model_label(model: str | None = None) -> str:
    """Human label for the model that will actually run (after availability fallback),
    e.g. for the report's 'done' frame."""
    return MODELS[_available_model(model)]["label"]

# Per-task ZeroGPU windows. A summary is short; a report is a long map-reduce, so
# it gets a bigger window. Bigger windows reserve more of the daily GPU-second
# quota per call, so keep summaries tight.
GPU_DURATION = int(os.getenv("GPU_DURATION", "90"))
SUMMARIZE_DURATION = int(os.getenv("GPU_DURATION_SUMMARIZE", str(GPU_DURATION)))
REPORT_DURATION = int(os.getenv("GPU_DURATION_REPORT", "180"))
# A single-pass item answer is one streamed call, so it needs a much smaller window
# than the map-reduce report. The report window still backs the map-reduce fallback.
ANSWER_DURATION = int(os.getenv("GPU_DURATION_ANSWER", "90"))

# Coalesce streamed answer tokens into ~this many characters per frame, so a long
# answer sends tens of progress frames to the browser instead of one per token.
_STREAM_FLUSH_CHARS = int(os.getenv("ANSWER_STREAM_FLUSH_CHARS", "32"))

# An agent turn is several model calls plus tool I/O (and possibly a nested
# summarize/report) inside ONE window — give it the widest budget. AGENT_MAX_STEPS
# caps tool calls so the turn can't run unbounded against that window.
AGENT_DURATION = int(os.getenv("GPU_DURATION_AGENT", "240"))
AGENT_MAX_STEPS = int(os.getenv("AGENT_MAX_STEPS", "6"))

# Cap how many chunks a report mines so the map-reduce fits REPORT_DURATION on
# ZeroGPU. generate_report surfaces the cap in the report ("covers the first N
# of M chunks"), so it's not silent.
REPORT_MAX_SECTIONS = int(os.getenv("REPORT_MAX_SECTIONS", "30"))

# Back-compat label (default model) — backend uses model_label(model) for the chosen one.
MODEL_LABEL = MODELS[DEFAULT_MODEL]["label"]


@lru_cache(maxsize=len(MODELS))
def model_path(model: str = DEFAULT_MODEL) -> str:
    """Resolve (downloading if needed) the local path to a model's GGUF file.

    Cached per model so each download happens once. Call :func:`prefetch_models`
    at startup so the first user request doesn't pay the download cost.
    """
    from huggingface_hub import hf_hub_download

    cfg = MODELS[resolve_model(model)]
    return hf_hub_download(cfg["repo"], cfg["file"], token=os.getenv("HF_TOKEN"))


def prefetch_models() -> dict[str, str]:
    """Download every selectable GGUF ahead of the first request (Space startup).

    Both models live on the Space so the request-time toggle never stalls on a
    download. Returns ``{model_key: local_path}``. A failure on one model (e.g. the
    private 26B without HF_TOKEN) is logged and skipped, not fatal.
    """
    paths: dict[str, str] = {}
    for key in MODELS:
        try:
            paths[key] = model_path(key)
        except Exception as e:  # noqa: BLE001
            print(f"[startup] could not prefetch model '{key}' "
                  f"({MODELS[key]['repo']}): {type(e).__name__}: {e}", flush=True)
    return paths


# Back-compat alias: older callers (and tests) may import prefetch_model.
def prefetch_model() -> str:
    """Prefetch just the default model's GGUF (kept for back-compat)."""
    return model_path(DEFAULT_MODEL)


def _load_llama(model: str = DEFAULT_MODEL):
    """Instantiate the chosen llama.cpp model. MUST be called inside a @spaces.GPU
    window so ``n_gpu_layers=-1`` actually binds to the allocated GPU."""
    from llama_cpp import Llama

    cfg = MODELS[resolve_model(model)]
    n_ctx = cfg.get("n_ctx") or N_CTX  # per-model cap (VRAM), else the global window
    return Llama(
        model_path=model_path(model),
        n_ctx=n_ctx,
        n_gpu_layers=N_GPU_LAYERS,
        verbose=False,
    )


# Which model GGUFs are actually fetchable in this process — memoized so a private
# model that 401s without HF_TOKEN (e.g. the 26B) is probed once, not on every call.
_AVAIL_CACHE: dict[str, bool] = {}


def _model_available(key: str) -> bool:
    """Whether a model's GGUF can be downloaded here (cached). The private 26B 401s
    when HF_TOKEN is unset/lacks access — we don't want that to crash a request."""
    if key in _AVAIL_CACHE:
        return _AVAIL_CACHE[key]
    try:
        model_path(key)
        ok = True
    except Exception:  # noqa: BLE001 - unavailable (e.g. private repo, no token)
        ok = False
    _AVAIL_CACHE[key] = ok
    return ok


def _available_model(model: str | None) -> str:
    """Resolve a request's model to one whose GGUF is actually fetchable.

    Falls back to the default (then any available) model when the requested one can't
    be downloaded — so picking the 26B on a Space without HF_TOKEN gracefully runs the
    E4B instead of raising a 401 mid-turn."""
    key = resolve_model(model)
    if _model_available(key):
        return key
    for k in (DEFAULT_MODEL, *MODELS):
        if k != key and _model_available(k):
            return k
    return key  # nothing fetchable — let the loader surface the real error


def available_models() -> list[str]:
    """Model keys whose GGUFs are fetchable here, in MODELS order (for the UI toggle)."""
    return [k for k in MODELS if _model_available(k)]


def _make_completer(llm):
    """A ``complete(prompt, system=None, max_tokens=...)`` callable over a llama
    instance, matching the signature chroma's summarize/report expect."""

    def complete(prompt: str, system: str | None = None,
                 max_tokens: int = MAX_OUTPUT_TOKENS) -> str:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        out = llm.create_chat_completion(
            messages=messages, temperature=0.2, max_tokens=max_tokens,
        )
        return (out["choices"][0]["message"]["content"] or "").strip()

    return complete


def _stream_chat(llm, prompt: str, system: str, max_tokens: int) -> Iterator[str]:
    """Yield answer text deltas from a streamed llama.cpp chat completion."""
    messages = [{"role": "system", "content": system},
                {"role": "user", "content": prompt}]
    for chunk in llm.create_chat_completion(
        messages=messages, temperature=0.2, max_tokens=max_tokens, stream=True,
    ):
        delta = (chunk["choices"][0].get("delta") or {}).get("content")
        if delta:
            yield delta


@spaces.GPU(duration=SUMMARIZE_DURATION)
def gpu_summarize(text: str, model: str = "") -> str:
    """Load the chosen model on the GPU and summarize ``text`` (whole map-reduce in
    one GPU window). Returns the markdown summary."""
    llm = _load_llama(_available_model(model))
    return summarize_text(text, complete=_make_completer(llm))


def _report_frames(llm, documents: list[dict], question: str,
                   max_sections: int) -> Iterator[dict]:
    """Run the map-reduce report over an already-loaded ``llm`` and stream frames.

    A worker thread runs the map-reduce while we relay its ``progress`` callback —
    threads share the process (and thus the allocated GPU); subprocesses would not.
    Shared by :func:`gpu_report` and :func:`gpu_answer`'s fallback so neither nests a
    second ``@spaces.GPU`` window.
    """
    max_sections = min(int(max_sections), REPORT_MAX_SECTIONS)
    complete = _make_completer(llm)
    q: "queue.Queue[tuple]" = queue.Queue()
    result: dict = {}

    def cb(frac: float, message: str) -> None:
        q.put(("progress", frac, message))

    def worker() -> None:
        try:
            result["report"] = generate_report(
                documents, question, complete=complete,
                max_chunks=int(max_sections), progress=cb,
            )
            q.put(("done", None))
        except Exception as e:  # noqa: BLE001
            q.put(("error", f"{type(e).__name__}: {e}"))

    threading.Thread(target=worker, daemon=True).start()

    while True:
        kind, *rest = q.get()
        if kind == "progress":
            yield {"frac": float(rest[0]), "message": rest[1]}
        elif kind == "done":
            yield {"report": result.get("report", ""), "done": True}
            return
        else:  # error
            yield {"error": rest[0]}
            return


@spaces.GPU(duration=REPORT_DURATION)
def gpu_report(documents: list[dict], question: str, max_sections: int,
               model: str = "") -> Iterator[dict]:
    """Stream a query-framed report, GPU held for the generator's lifetime.

    Yields ``{"frac", "message"}`` progress frames and a final
    ``{"report", "done": True}`` (or ``{"error"}``).
    """
    yield from _report_frames(_load_llama(_available_model(model)), documents, question,
                              max_sections)


@spaces.GPU(duration=ANSWER_DURATION)
def gpu_answer(documents: list[dict], question: str, max_sections: int,
               engine: str = "auto", model: str = "") -> Iterator[dict]:
    """Stream a single-pass answer about one agenda item, GPU held for the call.

    Assembles the most relevant packet context once and streams the answer token by
    token, yielding ``{"frac","message"}`` progress frames, then growing
    ``{"report": partial}`` frames, then a final ``{"report": full, "done": True}``.
    ``engine`` is the user's choice: ``"single"`` (semantic search), ``"mapreduce"``
    (read every section), or ``"auto"``; ``model`` selects which GGUF to load. Map-reduce
    reuses the same loaded model — no nested GPU window.
    """
    char_budget, max_tokens = budget_for(max_sections)
    llm = _load_llama(_available_model(model))

    if resolve_engine(documents, engine, char_budget=char_budget) == "mapreduce":
        yield from _report_frames(llm, documents, question, max_sections)
        return

    yield {"frac": 0.1, "message": "Finding the most relevant pages…"}
    try:
        context, info = assemble_context(documents, question, char_budget=char_budget)
        if not context.strip():
            yield {"report": "_No extractable text found for this item._", "done": True}
            return
        prompt = ANSWER_PROMPT.format(title=info["title"], query=question, context=context)
        yield {"frac": 0.35, "message": "Writing the answer…"}
        parts: list[str] = []
        emitted = 0  # chars already pushed to the client; coalesce tokens between frames
        for delta in _stream_chat(llm, prompt, ANSWER_SYSTEM, max_tokens):
            parts.append(delta)
            text = "".join(parts)
            if len(text) - emitted >= _STREAM_FLUSH_CHARS:
                emitted = len(text)
                yield {"frac": 0.6, "message": "Writing the answer…", "report": text}
        yield {"report": "".join(parts).strip(), "done": True}
    except Exception as e:  # noqa: BLE001
        yield {"error": f"{type(e).__name__}: {e}"}


@spaces.GPU(duration=AGENT_DURATION)
def gpu_agent_turn(
    upload_id: str, messages: list[dict], model: str = "",
) -> Iterator[dict]:
    """Run one Agent Mode turn inside a single GPU window.

    Loads the chosen GGUF once and hands the *same* completer to both the ReAct loop
    (its reasoning) and the tool context (so LLM-backed tools like summarize/report
    reuse this loaded model rather than trying to open a nested ``@spaces.GPU``
    window). Yields the loop's frames straight through (see
    :func:`webapp.agent_loop.agent_turn`).
    """
    from webapp.agent_loop import agent_turn
    from webapp.agent_tools import ToolContext

    llm = _load_llama(_available_model(model))
    complete = _make_completer(llm)
    ctx = ToolContext(upload_id=upload_id, complete=complete)
    yield from agent_turn(messages, ctx, complete, max_steps=AGENT_MAX_STEPS)


@spaces.GPU(duration=AGENT_DURATION)
def gpu_lii_agent_turn(messages: list[dict], model: str = "") -> Iterator[dict]:
    """Run one Cornell LII legal-research turn inside a single GPU window.

    Parallel to :func:`gpu_agent_turn` but with the LII toolkit and no uploaded packet —
    the tools query the eCFR API / Cornell LII over the network rather than the GPU, so
    the loaded GGUF is used only for the agent's reasoning. Yields the loop's frames
    straight through (see :func:`webapp.agent_loop.agent_turn`).
    """
    from webapp.agent_loop import agent_turn
    from webapp.lii_tools import TOOLKIT, LiiContext

    llm = _load_llama(_available_model(model))
    complete = _make_completer(llm)
    ctx = LiiContext()
    yield from agent_turn(messages, ctx, complete, max_steps=AGENT_MAX_STEPS, toolkit=TOOLKIT)