Files
stack/docs/superpowers/specs/2026-07-16-llm-module-design.md

11 KiB
Raw Blame History

llm module — local RAG over rulemaking comments + Zotero corpus

Status: Designed — awaiting implementation (milestones P33P35). Superseded in part by 2026-09-03-llm-corpus-recency-links-gpu-design.md (retrieval scope, ranking, evidence links, GPU routing). Date: 2026-07-16 Related: #253 (extraction, closed), #254/#255/#256 (P30, open — consumers of this module), .claude/specs/2026-04-23-comments-extraction-design.md (distributed-GPU notes)

Goal

A new llm module (src/llm, stack[llm] extra) that manages LangChain, a FastAPI service, and a vector DB to embed and query the regulations.gov public comments (164,214 items in the bib store, doctype:comment) and to tag them via RAG against the Zotero/bib reference corpus (rules, IOM, OIG, manuals). All inference is local — no cloud LLM APIs.

Decisions

Decision Choice Why
Inference server Ollama (new compose service, nvidia runtime) Model pull/quantize management, OpenAI-compatible API, idle unload (OLLAMA_KEEP_ALIVE) plays nice with the shared 3060, first-class langchain-ollama support. vLLM pins VRAM permanently; in-process llama-cpp couples API uptime to VRAM.
Vector DB pgvector on the existing Postgres No 30th stateful service; real concurrent writes (avoids repeating the DuckDB single-writer pain, #508#514); langchain-postgres; rides existing backups.
API surface Own FastAPI service at llm.fhirworx.io Keeps heavy LangChain deps out of the lean api image (skinny-install philosophy); tagging restarts don't touch the main API.
Tag vocabulary Closed vocab → bib Model picks from a curated set derived from existing corpus tags; written to bib Store under llm: prefix; nightly zotero-sync pushes to Zotero. Prefix separates machine tags from hand tags and lets sync scope with --tag.
Models Bake-off issue, not hardcoded Embeddings: nomic-embed-text vs bge-m3. Generation: 8B-class instruct Q4 as the floor (fits the 3060); larger models allowed on the 4090/5080 if the bake-off shows quality wins.
v1 consumers marimo notebook helper, stack llm CLI, minimal web search UI All three requested.

Architecture

bib.sqlite (164k comments + .md extractions)   Zotero/bib corpus (rules, IOM, OIG)
        \                                          /
         v                                        v
   src/llm/chunk.py  ── deterministic chunk ids (item key + content hash + seq)
        |
        v
   src/llm/index.py ── batch runner ──> Ollama /api/embed ──> pgvector (postgres, `llm` schema)
                          |                (1..N hosts)
                          v
   src/llm/rag.py  ── LangChain retriever + grounded generation (citations = bib item keys)
        |
        v
   src/llm/api.py  ── FastAPI: /health /search /similar /query /tag  @ llm.fhirworx.io
        |
        +── src/llm/tag.py ── closed-vocab RAG tagging ──> bib Store (`llm:` tags) ──> zotero-sync

Two new compose services: ollama (nvidia runtime, model volume, healthcheck) and llm (FastAPI behind Traefik/oauth2-proxy, OTel-wired like api).

Data flow

  • Index — chunker reads comment text from the per-attachment .md siblings written by the #253 extraction pipeline (abstract fallback when no attachment), plus corpus documents. Embeds via Ollama, upserts to pgvector. Incremental and resumable: keyed on bib item key + content hash; re-runs only embed new/changed items. Scopable: stack llm index --docket ….
  • Query — similarity search with metadata filters (docket/rule/year/ doctype) and optional grounded generation that answers only from retrieved context, citing bib item keys.
  • Tag — per comment: retrieve top-k corpus passages → closed-vocab structured output with confidence → write llm:<tag> to bib. Idempotent: re-tagging replaces prior llm: tags, never touches hand tags.

Multi-GPU strategy

Three machines: rack 3060 12 GB (shared with notebooks/Zotero), rig 4090, laptop 5080. The April spec sketched three fan-out options and leaned toward the simplest until inference workloads justified more — tagging 164k comments (~4 GPU-days at ~2 s/comment on one card) is that workload.

Design: remote GPUs are just Ollama endpoints. The rig and laptop run ollama serve; the batch runners accept LLM_OLLAMA_HOSTS (list) and dispatch chunks/comments least-loaded across hosts over LAN/tailscale. All data I/O (bib reads, pgvector/bib writes) stays on the server — remote boxes contribute pure compute, no shared FS, no queue service, no worker deploys. Per-host model availability is checked at startup; a missing host degrades gracefully to the remaining pool. Docket partitioning (--docket) remains the manual fallback. This is the "option 1-lite" the April spec anticipated, minus the task server it feared.

Scale

  • Embedding 164k comments (≈0.51M chunks × 768-dim): hours on one GPU, less fanned out; pgvector HNSW at this size is a few GB — fine.
  • Generative tagging is the bottleneck: batch, resumable, docket-scoped, off-peak on the rack card, fan out to 4090/5080 when available.
  • v1 gate: one docket end-to-end (index → query → tag → eval) before fan-out.

P33 build outcomes (2026-07-17)

P33 landed on branch llm-p33. Two infrastructure findings changed the plan:

pgvector ANN index build crashed the shared Postgres (AVX-512 SIGILL) — fixed (#580). The mirrored fhirworx/postgresql:latest image (Bitnami PG18, shared by gitea/nessie/polaris) shipped a pgvector 0.8.1 vector.so built with -march=native on an AVX-512 build host, but the server is a Ryzen 9 3950X (Zen 2, no AVX-512). CREATE INDEX … USING hnsw and ivfflat faulted with signal 4 (SIGILL) — even on 3 rows — because vector_norm() executed EVEX-encoded AVX-512 (vextractf64x2, vcvtusi2sd); the fault killed every backend and dropped the whole server (and its other databases) into recovery. Query-time distance ops were compiled as legal AVX2, which is why exact search never faulted. Fixed by infra/images/postgresql.Dockerfile: a patched fhirworx/postgresql:pgvector image built FROM the mirror that rebuilds pgvector with -march=x86-64-v3 (AVX2/FMA, fully supported by Zen 2) and swaps in the AVX-512-free vector.so. Verified: HNSW builds on the 15,852-vector pilot in ~2 s with no crash, queries use Index Scan using ix_embedding_hnsw, and dependents are unaffected. [llm].build_ann_index now defaults true; set it false to fall back to exact search. compose.yml and deploy.sh build the patched image from the mirror base.

Embed model bake-off (#564). Retrieval-proxy metric (comment abstract as query against its own pooled chunks, 120 pilot comments):

model dims recall@5 MRR chunks/s (1× 3060)
nomic-embed-text 768 86.7% 0.850 105
bge-m3 1024 93.3% 0.893 18

bge-m3 retrieves better (+6.6 pp recall@5) but is ~6× slower — ~25 h vs ~4 h to embed the full 164k-comment corpus on one card. P33 default stays nomic-embed-text (768-dim): the pilot is validated on it, 86.7% recall is solid, and it keeps the delivered index internally consistent. Recommended upgrade for the corpus-wide run: bge-m3, adopted at P35/scale-up time when the 4090/5080 pool absorbs the throughput cost — switching means embed_dim = 1024 and a --force re-index. Instruct-model selection is deferred to P35, where it is judged against the real closed vocab rather than a toy prompt.

Pilot result: docket CMS-2017-0092 indexed end-to-end — 1,618 comments → 15,852 chunks; a killed run resumed with 25 already-done items skipped (resumability proven); semantic queries return on-topic comments (telehealth originating-site §1834(m) exemptions; documentation-burden reduction). The v1 gate (one docket, index → query) is met; tag → eval arrives in P35.

Relationship to P30

#254 (position classification + thematic tagging) gets implemented on this module's tagging engine rather than duplicated; #255/#256 become downstream consumers of the query API. #416 (author/org metadata extraction) is adjacent but out of scope here.

Milestones and issues

P33: LLM Foundation — Ollama, pgvector, embedding index

  1. module scaffold + stack[llm] extrasrc/llm package; pyproject extra (langchain-core, langchain-ollama, langchain-postgres, fastapi, uvicorn, psycopg[binary], pydantic); [llm] section in stack.toml via conf; tests scaffold.
  2. ollama compose service + model bake-off — nvidia runtime, model volume, healthcheck, OLLAMA_KEEP_ALIVE idle unload; bake off embedding models (nomic-embed-text vs bge-m3) and instruct models with per-host sizing (3060 12 GB / 4090 24 GB / 5080 16 GB); record decision here.
  3. pgvector schema + migrations — enable extension, llm schema, comments + corpus collections, HNSW index, migration script, confirm backup ride-along.
  4. chunker — markdown-aware chunking of extraction .md files with abstract fallback; corpus doc chunking; deterministic chunk ids.
  5. incremental embedding indexer — batch runner bib→Ollama→pgvector, content-hash resumable, --docket/--collection scoping.
  6. multi-host Ollama dispatchLLM_OLLAMA_HOSTS fan-out, least-loaded dispatch, per-host model check, graceful degradation to local-only.

P34: LLM Query — RAG API, CLI, notebook client

  1. FastAPI service + compose + Traefik routellm.fhirworx.io behind oauth2-proxy; /health; OTel wiring like api.
  2. retrieval endpoints — /search (similarity + metadata filters), /similar/{key}, pagination, results cite bib item keys.
  3. grounded generation endpoint — /query: retrieve → generate with citations, streaming, answers only from retrieved context.
  4. stack llm CLI — typer subcommand: index / search / query / tag.
  5. marimo notebook helper — client module + example notebook (lake/OPPS pattern).

P35: LLM Tagging — closed-vocab RAG tagging into bib/Zotero

  1. tag vocabulary curation — derive closed vocab from existing corpus tags, curate, version it.
  2. RAG tagging chain + batch runner — top-k corpus retrieval → closed-vocab structured output with confidence; batch, resumable, docket-scoped, multi-host; implements #254's thematic tagging.
  3. write-back to bib + zotero-sync integrationllm: prefix, idempotent replace of prior llm: tags, hand tags untouched, scoped sync verified.
  4. golden-set eval — hand-labeled sample (~200 comments across dockets), precision/recall gate before corpus-wide fan-out; harness runs mocked in CI, real on demand.
  5. search web UI — minimal page at llm.fhirworx.io: search box, filters, results with tags + citations.
  6. observability — Prometheus metrics (embed/tag throughput, per-host dispatch, queue depth), Grafana panel; nvidia-exporter already covers GPU.

Testing

Unit tests mock Ollama and pgvector (coverage bar holds); one integration smoke test per milestone behind a marker; golden-set eval is the tagging quality gate.

Out of scope

  • Sentiment/coordination analytics (#255, #256) — consumers, not part of llm.
  • Author/org metadata extraction (#416).
  • OCR of image-only PDFs (phase-2 of the extraction spec).
  • Fine-tuning; clustering frameworks (Ray/Dask) — multi-host dispatch covers the need at this volume.