stack/src/llm/api.py

"""FastAPI app for the llm chat UI (P34).

Served at llm.fhirworx.io behind Traefik's git-sso middleware — every request
that reaches this app is already authenticated by Gitea SSO, so the app only
reads the forwarded identity header for display, it does not re-authenticate.

Run: ``uvicorn llm.api:app`` (or ``stack llm serve``).
"""

from __future__ import annotations

import json
import logging
import re
import time
from contextlib import asynccontextmanager
from importlib import resources
from pathlib import Path
from typing import AsyncIterator, Iterator

from fastapi import FastAPI, Header, HTTPException
from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse
from pydantic import BaseModel

log = logging.getLogger(__name__)


@asynccontextmanager
async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Warm the DuckDB replica (and refresh the derived families from
    it) at boot instead of on the chat's first code-valuation lookup —
    refs #699 ruling B3. Never raises: ``llm.evidence.warm`` is a quiet
    no-op when the replica file isn't there yet."""
    from llm import config as llm_config
    from llm.evidence import warm

    warm(llm_config.load())
    yield


app = FastAPI(
    title="llm chat",
    description=(
        "RAG chat over the CMS rulemaking library — comments, FR rules, "
        "reference corpus."
    ),
    version="0.2.0",
    lifespan=_lifespan,
)

_ISO_DATE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
_MODES = {"auto", "timeline", "recent"}
#: /search, /similar collection + kind query values — mirrors
#: llm.search.COLLECTIONS/KINDS (kept as plain literals here, like
#: _MODES above, rather than importing llm.search at module scope).
_COLLECTION_CHOICES = {"all", "comments", "rules", "corpus"}
_KIND_CHOICES = {"comment", "rule", "corpus"}

# OTel instrumentation — no-op if perf not installed or telemetry disabled.
try:
    from perf import init as _perf_init
    from perf.middleware import instrument as _perf_instrument

    _perf_init()
    _perf_instrument(app)
except ImportError:  # pragma: no cover — perf always installed in the image
    pass


class ChatRequest(BaseModel):
    question: str
    since: str | None = None
    #: "auto" (default — era-balanced retrieval for history-shaped
    #: questions, otherwise recency-blended), "timeline" or "recent" —
    #: forced from the UI (Task 6) and the golden eval (Task 7).
    #: Validated by ``chat()`` (400, not pydantic's 422) against
    #: ``_MODES``; the type stays ``str`` so an invalid value reaches
    #: that check instead of failing model validation first.
    mode: str = "auto"


def _page(name: str = "chat.html") -> str:
    return resources.files("llm").joinpath(f"web/{name}").read_text()


@app.get("/health")
def health() -> dict:
    """Liveness probe for the container healthcheck."""
    return {"status": "ok"}


@app.get("/", response_class=HTMLResponse)
def index() -> str:
    return _page()


@app.get("/ui/search", response_class=HTMLResponse)
@app.get("/ui/search/", response_class=HTMLResponse)
def search_page() -> str:
    """The search UI (refs #578) — a framework-free page like the chat,
    calling ``/search`` and ``/similar/{key}`` from the browser.

    Registered with and without the trailing slash so
    ``dev/scripts/nb_fe_smoke.py --url http://llm:8000/ui/search`` gets a
    200 from its readiness probe (which appends ``/``) instead of the
    307 Starlette would otherwise redirect with.
    """
    return _page("search.html")


_VENDOR = {"pdf.min.mjs", "pdf.worker.min.mjs", "jszip.min.js", "docx-preview.min.js"}


@app.get("/ui/vendor/{name}")
def vendor(name: str) -> Response:
    """Self-hosted browser libraries (PDF.js, JSZip + docx-preview) — a fixed allow-list, served
    from the package so the viewer works without a CDN and its worker is
    same-origin (a cross-origin worker script is refused by browsers)."""
    if name not in _VENDOR:
        raise HTTPException(404)
    body = resources.files("llm").joinpath(f"web/vendor/{name}").read_bytes()
    return Response(
        body,
        media_type="text/javascript",
        headers={"Cache-Control": "public, max-age=86400"},
    )


@app.get("/ui/pdf/{key}", response_class=HTMLResponse)
def pdf_page(key: str) -> str:
    """The PDF viewer; the page reads its item key, file, page and search
    text from its own URL and fetches /pdf/<key> for the file list."""
    if not _KEY_RE.match(key):
        raise HTTPException(404)
    return _page("pdf.html")


_KEY_RE = re.compile(r"^[A-Za-z0-9]{6,12}$")


def _pdf_paths() -> tuple[Path, Path, Path]:
    """(zotero.sqlite, zotero storage, comments root) — the comments root
    is the indexer's (``llm.source._default_root``): ``.state/comments``."""
    from conf import ROOT
    from conf import path as cpath

    return (
        Path(cpath("db.zotero")),
        Path(cpath("storage.zotero")),
        Path(ROOT) / ".state" / "comments",
    )


def _pdf_files(key: str) -> list:
    from conf.connect import bib
    from llm.pdfs import pdf_files

    store = bib()
    try:
        zdb, zstore, croot = _pdf_paths()
        return pdf_files(
            store, key, zotero_sqlite=zdb, zotero_storage=zstore, comments_root=croot
        )
    finally:
        store.close()


@app.get("/pdf/{key}")
def pdf_list(key: str) -> dict:
    """The attachments on file for a comment or library item: name, size,
    where it came from (comment download, bib attachment, Zotero storage)
    and how the viewer renders it (``pdf`` / ``docx`` / ``text``, or null
    for download only). Rules are not served here — they link to
    federalregister.gov."""
    if not _KEY_RE.match(key):
        raise HTTPException(404)
    files = _pdf_files(key)
    title = ""
    try:
        from conf.connect import bib

        store = bib()
        try:
            title = store.get(key).title or ""
        finally:
            store.close()
    except Exception:  # noqa: BLE001 — a Zotero-only key has no bib row
        title = ""
    return {
        "key": key,
        "title": title,
        "files": [
            {
                "name": f.name,
                "size": f.size,
                "source": f.source,
                "renderable": f.renderable,
                "renderer": f.renderer,
                "media_type": f.media_type,
            }
            for f in files
        ],
    }


@app.get("/pdf/{key}/file")
def pdf_file(key: str, name: str = "") -> FileResponse:
    """The attachment bytes — a PDF inline with range requests (PDF.js
    streams pages from a large file instead of waiting for all of it),
    anything else as a download."""
    from llm.pdfs import resolve

    if not _KEY_RE.match(key):
        raise HTTPException(404)
    files = _pdf_files(key)
    from conf.connect import bib

    store = bib()
    try:
        bib_root = Path(getattr(store, "_storage", "storage"))
    finally:
        store.close()
    _zdb, zstore, croot = _pdf_paths()
    chosen = resolve(files, name, roots=(bib_root, zstore, croot))
    if chosen is None:
        raise HTTPException(404, "no such attachment on file for this item")
    return FileResponse(
        chosen.path,
        media_type=chosen.media_type,
        filename=chosen.name,
        content_disposition_type="inline" if chosen.renderer == "pdf" else "attachment",
        headers={"Accept-Ranges": "bytes", "Cache-Control": "private, max-age=3600"},
    )


@app.get("/whoami")
def whoami(x_auth_request_user: str = Header(default="")) -> dict:
    """The Gitea username Traefik forwarded (empty if unset)."""
    return {"user": x_auth_request_user}


def _sse(question: str, since: str = "", mode: str = "auto") -> Iterator[str]:
    from llm import config as llm_config
    from llm.pool import HostPool
    from llm.rag import stream_answer

    cfg = llm_config.load()
    pool = HostPool.from_config(cfg)
    try:
        for event in stream_answer(
            question, cfg=cfg, pool=pool, since=since, mode=mode
        ):
            yield f"data: {json.dumps(event)}\n\n"
    except Exception as exc:  # surface to the transcript, don't 500 mid-stream
        yield f"data: {json.dumps({'type': 'error', 'message': str(exc)})}\n\n"


@app.post("/chat")
def chat(req: ChatRequest) -> StreamingResponse:
    question = req.question.strip()
    if not question:
        raise HTTPException(status_code=400, detail="empty question")
    since = (req.since or "").strip()
    if since and not _ISO_DATE.match(since):
        raise HTTPException(status_code=400, detail="since must be YYYY-MM-DD")
    mode = (req.mode or "auto").strip()
    if mode not in _MODES:
        raise HTTPException(
            status_code=400, detail="mode must be one of: auto, timeline, recent"
        )
    return StreamingResponse(
        _sse(question, since, mode), media_type="text/event-stream"
    )


@app.get("/hosts")
def hosts() -> dict:
    """The Ollama fleet as the service sees it right now: declared VRAM,
    liveness, models, and the host + model a chat would use."""
    from llm import config as llm_config
    from llm.pool import HostPool, pick_model

    cfg = llm_config.load()
    pool = HostPool.from_config(cfg)
    declared = [{"host": h, "vram_gb": pool.vram(h)} for h in pool.hosts]
    try:
        live = set(pool.check(cfg.instruct_model))
    except RuntimeError as exc:
        return {
            "hosts": [{**d, "live": False, "models": []} for d in declared],
            "generation": None,
            "error": str(exc),
        }
    status = {row["host"]: row for row in pool.status()}
    rows = [
        {
            **d,
            "live": d["host"] in live,
            "models": status.get(d["host"], {}).get("models", []),
        }
        for d in declared
    ]
    with pool.acquire_generation() as host:
        model = pick_model(cfg, pool, host)
    return {"hosts": rows, "generation": {"host": host, "model": model}}


def _limit_ok(limit: int) -> bool:
    return 1 <= limit <= 50


@app.get("/search")
def search_endpoint(
    q: str = "",
    collection: str = "all",
    docket: str = "",
    item_key: str = "",
    year: str = "",
    kind: str = "",
    limit: int = 10,
    offset: int = 0,
) -> dict:
    """Metadata-filtered similarity search over the indexed library.

    ``total`` is the number of results in this page (``len(results)``) —
    the underlying search overfetches per collection rather than running
    a separate COUNT query, so there is no cheaper true total to report.
    """
    from llm import config as llm_config
    from llm.pool import HostPool
    from llm.search import search as run_search

    query = q.strip()
    if not query:
        raise HTTPException(status_code=400, detail="q is required")
    if collection not in _COLLECTION_CHOICES:
        raise HTTPException(status_code=400, detail="unknown collection")
    if kind and kind not in _KIND_CHOICES:
        raise HTTPException(status_code=400, detail="unknown kind")
    if not _limit_ok(limit):
        raise HTTPException(status_code=400, detail="limit must be between 1 and 50")
    if offset < 0:
        raise HTTPException(status_code=400, detail="offset must be >= 0")

    filters = {
        k: v
        for k, v in {
            "docket": docket,
            "item_key": item_key,
            "year": year,
            "kind": kind,
        }.items()
        if v
    }
    cfg = llm_config.load()
    pool = HostPool.from_config(cfg)
    t0 = time.monotonic()
    results = run_search(
        query,
        cfg=cfg,
        pool=pool,
        collection=collection,
        filters=filters,
        limit=limit,
        offset=offset,
    )
    log.info(
        "search %r collection=%s filters=%s -> %d results in %.1fms",
        query,
        collection,
        filters,
        len(results),
        (time.monotonic() - t0) * 1000,
    )
    return {
        "query": query,
        "filters": filters,
        "total": len(results),
        "results": results,
    }


@app.get("/similar/{key}")
def similar_endpoint(key: str, collection: str = "all", limit: int = 10) -> dict:
    """Nearest neighbours of an already-indexed bib item's own chunk(s),
    excluding the item itself. 404 when *key* has no indexed chunks."""
    from llm import config as llm_config
    from llm.pool import HostPool
    from llm.search import similar as run_similar

    if collection not in _COLLECTION_CHOICES:
        raise HTTPException(status_code=400, detail="unknown collection")
    if not _limit_ok(limit):
        raise HTTPException(status_code=400, detail="limit must be between 1 and 50")

    cfg = llm_config.load()
    pool = HostPool.from_config(cfg)
    t0 = time.monotonic()
    results = run_similar(key, cfg=cfg, pool=pool, collection=collection, limit=limit)
    if results is None:
        raise HTTPException(status_code=404, detail=f"no indexed chunks for {key!r}")
    log.info(
        "similar %r collection=%s -> %d results in %.1fms",
        key,
        collection,
        len(results),
        (time.monotonic() - t0) * 1000,
    )
    return {"key": key, "total": len(results), "results": results}