60 KiB
LLM Foundation (P33) Implementation Plan
For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (
- [ ]) syntax for tracking.
Goal: Stand up the llm module's foundation — Ollama compose service, pgvector on the existing Postgres, chunker, multi-host embedding pool, and an incremental indexer over the regulations.gov comments and bib corpus.
Architecture: src/llm is a new skinny module (stack[llm] extra). Text sources (comment extraction .md files under .state/comments/, bib corpus attachments) → markdown-aware chunker with deterministic ids → embeddings via a least-loaded pool of Ollama hosts (rack 3060 + optional rig 4090 / laptop 5080) → langchain-postgres PGVector collections in a new llm database on the existing Bitnami Postgres (pgvector 0.8.1 confirmed available, vector.so present). Resume state lives in an index_state table keyed on (item_key, collection, content_hash).
Tech Stack: Python 3.12, httpx (Ollama /api/embed), langchain-postgres + psycopg3 + SQLAlchemy DSN, pgvector, typer CLI, pytest + unittest.mock.
Tracker: milestone P33 (id 34) — issues #563 (scaffold), #564 (ollama + bake-off), #565 (pgvector), #566 (chunker), #567 (indexer), #568 (multi-host dispatch). Spec: docs/superpowers/specs/2026-07-16-llm-module-design.md.
Global Constraints
- Coverage bar is 99% (
--cov-fail-under=99in CI); every new source file ships with tests in the same commit. - Ruff:
select = ["E", "F", "I"], line-length 88. Runuv run ruff check src/ tests/ && uv run ruff format src/ tests/before every commit. - Lazy conf imports:
from conf import ...only inside functions, never module top-level (repo convention; keeps optional extras importable). - Mocking: stdlib
unittest.mock+monkeypatchonly — no respx, no pytest-mock (not in dev deps). - Tests live in
tests/llm/mirroringsrc/llm/, with emptytests/llm/__init__.py. - New module must be added to
[tool.uv.build-backend] module-namein pyproject.toml or it won't package. .gitea/workflows/*.ymlare generated — never hand-edit; regen viauv run python dev/scripts/gen_config.pyif stack.toml[ci]inputs change (this plan doesn't change them).- Commit messages: conventional commits (
feat(llm): …), reference issues with(refs #N). No Co-Authored-By trailers. - Postgres is reachable from containers at
postgres:5432(storage net) and — after Task 2 — from the host at127.0.0.1:5432. Secrets come from.envviaos.environ, never stack.toml. - Ollama model names used throughout until the bake-off (Task 8) says otherwise: embed
nomic-embed-text(768-dim), instructllama3.1:8b.
Task 1: Module scaffold, stack[llm] extra, [llm] config (#563)
Files:
- Create:
src/llm/__init__.py,src/llm/config.py - Create:
tests/llm/__init__.py(empty),tests/llm/test_config.py - Modify:
pyproject.toml(extras + module-name list),stack.toml([llm],[services])
Interfaces:
-
Consumes:
conf.cfg(attribute access over stack.toml). -
Produces:
llm.config.LlmConfig(frozen dataclass) andllm.config.load() -> LlmConfig;llm.config.pg_url(cfg: LlmConfig) -> strreturning a SQLAlchemy/psycopg URL. All later tasks import these exact names. -
Step 1: Add config + packaging entries
stack.toml — add after the [bcda] section:
[llm]
ollama = "http://127.0.0.1:11434" # host-side default; containers override via LLM_OLLAMA_HOSTS
embed_model = "nomic-embed-text" # 768-dim; bake-off (P33 #564) may revise
instruct_model = "llama3.1:8b" # bake-off may revise
embed_dim = 768
pg_host = "127.0.0.1" # host-side default; containers set LLM_PG_HOST=postgres
pg_port = 5432
pg_db = "llm"
pg_user = "llm"
stack.toml — in [services], add:
ollama = "http://ollama:11434"
pyproject.toml — add to [project.optional-dependencies] (after bls):
llm = [
"stack[conf]",
"stack[bib]",
"httpx>=0.28.1",
"langchain-core>=0.3.0",
"langchain-ollama>=0.2.0",
"langchain-postgres>=0.0.12",
"psycopg[binary]>=3.2.0",
"sqlalchemy>=2.0.0",
]
Add "stack[llm]" to the cli extra list and to the all extra list. Add "llm" to [tool.uv.build-backend] module-name (keep the list alphabetized).
- Step 2: Write the failing test
tests/llm/test_config.py:
"""llm.config — [llm] section parsing + env overrides."""
import pytest
from llm import config as llm_config
class TestLoad:
def test_defaults_from_stack_toml(self, monkeypatch):
monkeypatch.delenv("LLM_OLLAMA_HOSTS", raising=False)
monkeypatch.delenv("LLM_PG_HOST", raising=False)
cfg = llm_config.load()
assert cfg.ollama_hosts == ("http://127.0.0.1:11434",)
assert cfg.embed_model == "nomic-embed-text"
assert cfg.embed_dim == 768
assert cfg.pg_host == "127.0.0.1"
assert cfg.pg_db == "llm"
def test_env_overrides(self, monkeypatch):
monkeypatch.setenv(
"LLM_OLLAMA_HOSTS",
"http://rig:11434, http://laptop:11434",
)
monkeypatch.setenv("LLM_PG_HOST", "postgres")
cfg = llm_config.load()
assert cfg.ollama_hosts == ("http://rig:11434", "http://laptop:11434")
assert cfg.pg_host == "postgres"
class TestPgUrl:
def test_url_includes_password_from_env(self, monkeypatch):
monkeypatch.setenv("LLM_DB_PASSWORD", "s3cret")
cfg = llm_config.load()
url = llm_config.pg_url(cfg)
assert url == f"postgresql+psycopg://llm:s3cret@{cfg.pg_host}:5432/llm"
def test_missing_password_raises(self, monkeypatch):
monkeypatch.delenv("LLM_DB_PASSWORD", raising=False)
cfg = llm_config.load()
with pytest.raises(RuntimeError, match="LLM_DB_PASSWORD"):
llm_config.pg_url(cfg)
- Step 3: Run test to verify it fails
Run: uv sync --extra llm && uv run pytest tests/llm/test_config.py -v
Expected: FAIL with ModuleNotFoundError: No module named 'llm'
- Step 4: Write the implementation
src/llm/__init__.py (bib-style: docstring + redundant-alias re-exports):
"""llm: local RAG over rulemaking comments + the Zotero/bib corpus.
Embeds regulations.gov public comments (``doctype:comment`` in the bib
store) and the reference corpus into pgvector, and (P34/P35) exposes
retrieval, grounded generation, and closed-vocab tagging — all inference
on local Ollama hosts, never cloud APIs.
Architecture::
.state/comments/*.md + bib corpus (llm.source)
-> markdown-aware chunks (llm.chunk)
-> Ollama /api/embed, multi-host (llm.pool)
-> pgvector collections (llm.index)
Spec: docs/superpowers/specs/2026-07-16-llm-module-design.md
"""
from llm.config import LlmConfig as LlmConfig
from llm.config import load as load
from llm.config import pg_url as pg_url
src/llm/config.py:
"""[llm] configuration — stack.toml section + env overrides.
Env contract:
LLM_OLLAMA_HOSTS comma-separated Ollama base URLs (beats [llm].ollama)
LLM_PG_HOST beats [llm].pg_host (containers set this to "postgres")
LLM_DB_PASSWORD required for pg_url(); lives in .env, never stack.toml
"""
from __future__ import annotations
import os
from dataclasses import dataclass
@dataclass(frozen=True)
class LlmConfig:
ollama_hosts: tuple[str, ...]
embed_model: str
instruct_model: str
embed_dim: int
pg_host: str
pg_port: int
pg_db: str
pg_user: str
def load() -> LlmConfig:
"""Read the [llm] section; env vars override host-ish values."""
from conf import cfg
section = cfg.llm
hosts_env = os.environ.get("LLM_OLLAMA_HOSTS", "")
if hosts_env:
hosts = tuple(h.strip() for h in hosts_env.split(",") if h.strip())
else:
hosts = (str(section.ollama),)
return LlmConfig(
ollama_hosts=hosts,
embed_model=str(section.embed_model),
instruct_model=str(section.instruct_model),
embed_dim=int(section.embed_dim),
pg_host=os.environ.get("LLM_PG_HOST", str(section.pg_host)),
pg_port=int(section.pg_port),
pg_db=str(section.pg_db),
pg_user=str(section.pg_user),
)
def pg_url(cfg: LlmConfig) -> str:
"""SQLAlchemy/psycopg URL for the llm database.
Password comes from LLM_DB_PASSWORD only — secrets never live in
stack.toml.
"""
password = os.environ.get("LLM_DB_PASSWORD", "")
if not password:
raise RuntimeError(
"LLM_DB_PASSWORD not set — add it to .env (see Task 2 provisioning)"
)
return (
f"postgresql+psycopg://{cfg.pg_user}:{password}"
f"@{cfg.pg_host}:{cfg.pg_port}/{cfg.pg_db}"
)
- Step 5: Run tests, lint, commit
Run: uv run pytest tests/llm/ -v → PASS (4 tests). uv run ruff check src/llm tests/llm && uv run ruff format src/llm tests/llm
git add src/llm tests/llm pyproject.toml stack.toml uv.lock
git commit -m "feat(llm): module scaffold, stack[llm] extra, [llm] config (refs #563)"
Task 2: Postgres provisioning — llm role/db, pgvector, host port (#565)
Files:
- Modify:
src/api/auth/manifest.py(POSTGRES_ROLES, POSTGRES_DATABASES, credential manifest) - Modify:
compose.yml(postgresports:) - Create:
src/llm/migrate.py,tests/llm/test_migrate.py - Modify:
tests/api/test_manifest.pyif role/db counts are asserted (check first:grep -rn "POSTGRES_ROLES\|POSTGRES_DATABASES" tests/)
Interfaces:
-
Consumes:
llm.config.load,llm.config.pg_url. -
Produces:
llm.migrate.migrate(engine) -> None(idempotent; createsindex_state),llm.migrate.ensure_hnsw(engine, dim: int) -> None,llm.migrate.INDEX_STATE_DDL: str. Task 7 (indexer) calls both. -
Step 1: Register the credential + role/db in the auth manifest
In src/api/auth/manifest.py, extend the two dicts (exact current values shown; add the llm lines):
POSTGRES_ROLES = {
"git": "GITEA_DB_PASSWORD",
"nessie": "NESSIE_DB_PASSWORD",
"polaris": "POLARIS_DB_PASSWORD",
"llm": "LLM_DB_PASSWORD",
}
POSTGRES_DATABASES = {
"git": "gitea",
"nessie": "nessie",
"polaris": "polaris",
"llm": "llm",
}
Also add an LLM_DB_PASSWORD credential entry following the exact pattern of the NESSIE_DB_PASSWORD entry in the manifest list (same tier/provisioner/format — copy that entry and change the env_var and description). Read the file first; match whatever dataclass/fields the neighboring entries use.
- Step 2: Expose postgres to the host (loopback only)
In compose.yml, add to the postgres service (after networks:):
ports:
# Loopback-only: host-side llm batch runs (indexer/tagger) need
# pgvector; everything else still uses the storage net.
- "127.0.0.1:5432:5432"
- Step 3: Write the failing migrate test
tests/llm/test_migrate.py — pure unit test: assert DDL content and that migrate() executes each statement on a mock engine (no live postgres in CI):
"""llm.migrate — idempotent DDL for the llm database."""
from unittest.mock import MagicMock
from llm import migrate
class TestDdl:
def test_index_state_ddl_is_idempotent_sql(self):
assert "CREATE TABLE IF NOT EXISTS index_state" in migrate.INDEX_STATE_DDL
assert "PRIMARY KEY (item_key, collection)" in migrate.INDEX_STATE_DDL
def test_migrate_executes_ddl(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.migrate(engine)
executed = " ".join(
str(call.args[0]) for call in conn.execute.call_args_list
)
assert "index_state" in executed
def test_ensure_hnsw_types_column_then_indexes(self):
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
migrate.ensure_hnsw(engine, 768)
executed = " ".join(
str(call.args[0]) for call in conn.execute.call_args_list
)
assert "vector(768)" in executed
assert "USING hnsw" in executed
- Step 4: Run test to verify it fails
Run: uv run pytest tests/llm/test_migrate.py -v
Expected: FAIL with ImportError: cannot import name 'migrate'
- Step 5: Write the implementation
src/llm/migrate.py:
"""Idempotent DDL for the llm database.
The ``vector`` extension and the ``llm`` role/database are provisioned by
the auth manifest (``api.auth.provision.bootstrap_postgres``) as superuser;
everything here runs as the unprivileged ``llm`` role.
``index_state`` is the resume ledger: one row per (item, collection) with
the content hash that was last embedded. The indexer skips items whose
hash is unchanged and re-embeds the rest.
"""
from __future__ import annotations
from sqlalchemy import text
from sqlalchemy.engine import Engine
INDEX_STATE_DDL = """
CREATE TABLE IF NOT EXISTS index_state (
item_key TEXT NOT NULL,
collection TEXT NOT NULL,
content_hash TEXT NOT NULL,
chunk_count INTEGER NOT NULL DEFAULT 0,
indexed_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (item_key, collection)
)
"""
# langchain-postgres creates langchain_pg_embedding with an untyped vector
# column; HNSW needs a typed one. ALTER is a no-op when already typed.
_HNSW_DDL = [
"ALTER TABLE langchain_pg_embedding "
"ALTER COLUMN embedding TYPE vector({dim})",
"CREATE INDEX IF NOT EXISTS ix_embedding_hnsw "
"ON langchain_pg_embedding USING hnsw (embedding vector_cosine_ops)",
]
def migrate(engine: Engine) -> None:
"""Create llm-owned tables. Safe to run on every start."""
with engine.begin() as conn:
conn.execute(text(INDEX_STATE_DDL))
def ensure_hnsw(engine: Engine, dim: int) -> None:
"""Type the embedding column and build the HNSW index.
Call after the PGVector store has created its tables (first
``add_embeddings``), not before.
"""
with engine.begin() as conn:
for ddl in _HNSW_DDL:
conn.execute(text(ddl.format(dim=dim)))
- Step 6: Run tests, lint
Run: uv run pytest tests/llm/ tests/api/ -q → PASS. Ruff check + format.
- Step 7: Provision the live database
# generate + store the password (provision flow owns .env writes; if the
# manifest provisioner requires the full bootstrap, use that instead):
uv run python -m api.auth bootstrap --only LLM_DB_PASSWORD $(git rev-parse HEAD) \
|| echo "LLM_DB_PASSWORD=$(openssl rand -hex 24)" >> .env
docker compose up -d postgres # picks up the new ports: mapping
# role + db (idempotent, mirrors bootstrap_postgres):
PW_SUPER=$(grep '^POSTGRES_PASSWORD=' .env | cut -d= -f2-)
PW_LLM=$(grep '^LLM_DB_PASSWORD=' .env | cut -d= -f2-)
docker exec -e PGPASSWORD="$PW_SUPER" postgres psql -U postgres <<SQL
DO \$\$ BEGIN CREATE ROLE llm LOGIN; EXCEPTION WHEN duplicate_object THEN NULL; END \$\$;
ALTER ROLE llm PASSWORD '$PW_LLM';
SELECT 'CREATE DATABASE llm OWNER llm' WHERE NOT EXISTS
(SELECT FROM pg_database WHERE datname = 'llm')\gexec
SQL
docker exec -e PGPASSWORD="$PW_SUPER" postgres psql -U postgres -d llm \
-c "CREATE EXTENSION IF NOT EXISTS vector;"
Note: check whether api.auth bootstrap supports --only; if not, append the openssl line to .env by hand and file nothing — the manifest entry keeps rotation covered going forward.
- Step 8: Verify from the host, then commit
PW_LLM=$(grep '^LLM_DB_PASSWORD=' .env | cut -d= -f2-)
uv run python -c "
import os; os.environ['LLM_DB_PASSWORD'] = '$PW_LLM'
import sqlalchemy, llm
from llm import migrate
eng = sqlalchemy.create_engine(llm.pg_url(llm.load()))
migrate.migrate(eng)
print(eng.connect().execute(sqlalchemy.text('SELECT 1')).scalar())
"
Expected output: 1.
git add src/api/auth/manifest.py compose.yml src/llm/migrate.py tests/llm/test_migrate.py
git commit -m "feat(llm): pgvector provisioning — llm role/db, loopback port, migrate (refs #565)"
Task 3: Ollama compose service (#564, infra half)
Files:
- Modify:
compose.yml(newollamaservice +ollama_modelsvolume)
Interfaces:
-
Produces: Ollama API at
http://127.0.0.1:11434(host/LAN) andhttp://ollama:11434(data net). Tasks 4, 7, 8 depend on it. -
Step 1: Add the service
In compose.yml, after the notebooks block (GPU convention: deploy.resources.reservations.devices, like notebooks — not the legacy runtime: nvidia):
ollama:
image: ollama/ollama:latest
container_name: ollama
networks:
- data
ports:
# Host + LAN: the rack's Ollama is one endpoint in the multi-host
# pool (rig 4090 / laptop 5080 serve the same API on the LAN).
- "11434:11434"
environment:
# Unload models after idle so the shared 3060 frees VRAM for
# notebooks/zotero.
- OLLAMA_KEEP_ALIVE=5m
volumes:
- ollama_models:/root/.ollama
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
healthcheck:
test: ["CMD", "ollama", "ls"]
interval: 30s
timeout: 10s
retries: 3
start_period: 30s
labels:
- "promtail=true"
security_opt:
- no-new-privileges:true
restart: unless-stopped
Add ollama_models: to the top-level volumes: block.
- Step 2: Start, pull models, verify
docker compose up -d ollama
docker exec ollama ollama pull nomic-embed-text
docker exec ollama ollama pull llama3.1:8b
curl -s http://127.0.0.1:11434/api/tags | python3 -m json.tool | grep name
curl -s http://127.0.0.1:11434/api/embed \
-d '{"model":"nomic-embed-text","input":["hello"]}' \
| python3 -c "import json,sys; e=json.load(sys.stdin)['embeddings'][0]; print(len(e))"
nvidia-smi --query-gpu=memory.used --format=csv # then again after ~5m idle
Expected: both models listed; embed returns 768; VRAM drops back after keep-alive expiry.
- Step 3: Commit
git add compose.yml
git commit -m "feat(llm): ollama compose service — GPU, idle unload, models volume (refs #564)"
Task 4: Multi-host Ollama pool (#568)
Files:
- Create:
src/llm/pool.py,tests/llm/test_pool.py - Modify:
src/llm/__init__.py(re-exportHostPool)
Interfaces:
-
Consumes:
llm.config.LlmConfig. -
Produces (Tasks 7–8 and P34/P35 consume these exact names):
class HostPoolwithHostPool(hosts: Sequence[str]),HostPool.from_config(cfg: LlmConfig) -> HostPool,check(model: str) -> list[str](drops hosts that are down or lack the model; raisesRuntimeErrorif none remain),acquire()context manager yielding the least-in-flight base URL.embed_texts(pool: HostPool, model: str, texts: Sequence[str], *, batch_size: int = 64) -> list[list[float]]— order-preserving, fans batches across hosts with threads.class PoolEmbeddings(Embeddings)— langchain adapter:PoolEmbeddings(pool, model),embed_documents,embed_query.
-
Step 1: Write the failing tests
tests/llm/test_pool.py:
"""llm.pool — least-loaded fan-out across Ollama hosts."""
import json
from unittest.mock import MagicMock, patch
import pytest
from llm.pool import HostPool, PoolEmbeddings, embed_texts
H1, H2 = "http://h1:11434", "http://h2:11434"
def _resp(payload, status=200):
r = MagicMock()
r.status_code = status
r.json.return_value = payload
return r
class TestCheck:
@patch("llm.pool.httpx.Client")
def test_drops_host_missing_model(self, MockClient):
client = MockClient.return_value.__enter__.return_value
client.get.side_effect = [
_resp({"models": [{"name": "nomic-embed-text:latest"}]}),
_resp({"models": [{"name": "other:latest"}]}),
]
pool = HostPool([H1, H2])
assert pool.check("nomic-embed-text") == [H1]
@patch("llm.pool.httpx.Client")
def test_drops_unreachable_host(self, MockClient):
import httpx
client = MockClient.return_value.__enter__.return_value
client.get.side_effect = [
_resp({"models": [{"name": "m:latest"}]}),
httpx.ConnectError("down"),
]
pool = HostPool([H1, H2])
assert pool.check("m") == [H1]
@patch("llm.pool.httpx.Client")
def test_no_hosts_left_raises(self, MockClient):
client = MockClient.return_value.__enter__.return_value
client.get.return_value = _resp({"models": []})
with pytest.raises(RuntimeError, match="no Ollama host"):
HostPool([H1]).check("m")
class TestAcquire:
def test_prefers_least_in_flight(self):
pool = HostPool([H1, H2])
with pool.acquire() as first, pool.acquire() as second:
assert {first, second} == {H1, H2}
def test_released_host_is_reused(self):
pool = HostPool([H1, H2])
with pool.acquire() as first:
pass
with pool.acquire() as a, pool.acquire() as b:
assert {a, b} == {H1, H2}
assert first in {H1, H2}
class TestEmbedTexts:
@patch("llm.pool.httpx.Client")
def test_order_preserved_across_batches(self, MockClient):
client = MockClient.return_value.__enter__.return_value
def fake_post(url, json=None, **kw):
return _resp(
{"embeddings": [[float(len(t))] for t in json["input"]]}
)
client.post.side_effect = fake_post
pool = HostPool([H1, H2])
out = embed_texts(pool, "m", ["a", "bb", "ccc"], batch_size=1)
assert out == [[1.0], [2.0], [3.0]]
@patch("llm.pool.httpx.Client")
def test_http_error_raises(self, MockClient):
client = MockClient.return_value.__enter__.return_value
resp = _resp({}, status=500)
resp.raise_for_status.side_effect = Exception("boom")
client.post.return_value = resp
with pytest.raises(Exception, match="boom"):
embed_texts(HostPool([H1]), "m", ["a"])
class TestPoolEmbeddings:
@patch("llm.pool.embed_texts")
def test_adapter_delegates(self, mock_embed):
mock_embed.return_value = [[0.1]]
pool = HostPool([H1])
emb = PoolEmbeddings(pool, "m")
assert emb.embed_documents(["x"]) == [[0.1]]
assert emb.embed_query("x") == [0.1]
- Step 2: Run tests to verify they fail
Run: uv run pytest tests/llm/test_pool.py -v
Expected: FAIL with ModuleNotFoundError: No module named 'llm.pool'
- Step 3: Write the implementation
src/llm/pool.py:
"""Least-loaded fan-out across Ollama hosts.
Remote GPUs are plain Ollama endpoints (rig 4090 / laptop 5080 run
``ollama serve``); this pool spreads embed calls across them while all
bib/pgvector I/O stays on the server. Uses the raw ``/api/embed`` HTTP API
rather than langchain-ollama so one pool can juggle N base URLs;
``PoolEmbeddings`` adapts it back to the langchain interface for PGVector.
"""
from __future__ import annotations
import threading
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from typing import Iterator, Sequence
import httpx
from langchain_core.embeddings import Embeddings
_TIMEOUT = httpx.Timeout(120.0, connect=5.0)
class HostPool:
"""Tracks in-flight requests per host; hands out the idlest one."""
def __init__(self, hosts: Sequence[str]) -> None:
self._lock = threading.Lock()
self._in_flight: dict[str, int] = {h.rstrip("/"): 0 for h in hosts}
@classmethod
def from_config(cls, cfg) -> "HostPool":
return cls(cfg.ollama_hosts)
@property
def hosts(self) -> list[str]:
return list(self._in_flight)
def check(self, model: str) -> list[str]:
"""Keep only hosts that are up and serve ``model``.
Ollama tags models ``name:latest``; match on the bare prefix.
"""
alive: list[str] = []
with httpx.Client(timeout=_TIMEOUT) as client:
for host in self.hosts:
try:
resp = client.get(f"{host}/api/tags")
names = {
m["name"].split(":")[0]
for m in resp.json().get("models", [])
}
if model.split(":")[0] in names:
alive.append(host)
except httpx.HTTPError:
continue
with self._lock:
self._in_flight = {h: 0 for h in alive}
if not alive:
raise RuntimeError(
f"no Ollama host in pool serves {model!r} — "
f"pull it or fix LLM_OLLAMA_HOSTS"
)
return alive
@contextmanager
def acquire(self) -> Iterator[str]:
with self._lock:
host = min(self._in_flight, key=self._in_flight.__getitem__)
self._in_flight[host] += 1
try:
yield host
finally:
with self._lock:
self._in_flight[host] -= 1
def embed_texts(
pool: HostPool,
model: str,
texts: Sequence[str],
*,
batch_size: int = 64,
) -> list[list[float]]:
"""Embed ``texts`` in order, fanning batches across the pool."""
batches = [
(i, list(texts[i : i + batch_size]))
for i in range(0, len(texts), batch_size)
]
results: dict[int, list[list[float]]] = {}
def run(start: int, batch: list[str]) -> None:
with pool.acquire() as host, httpx.Client(timeout=_TIMEOUT) as client:
resp = client.post(
f"{host}/api/embed", json={"model": model, "input": batch}
)
resp.raise_for_status()
results[start] = resp.json()["embeddings"]
workers = max(1, min(len(pool.hosts) * 2, len(batches)))
with ThreadPoolExecutor(max_workers=workers) as ex:
for future in [ex.submit(run, s, b) for s, b in batches]:
future.result()
return [vec for start in sorted(results) for vec in results[start]]
class PoolEmbeddings(Embeddings):
"""langchain adapter so PGVector can query through the pool."""
def __init__(self, pool: HostPool, model: str) -> None:
self._pool = pool
self._model = model
def embed_documents(self, texts: list[str]) -> list[list[float]]:
return embed_texts(self._pool, self._model, texts)
def embed_query(self, text: str) -> list[float]:
return self.embed_documents([text])[0]
Add to src/llm/__init__.py:
from llm.pool import HostPool as HostPool
from llm.pool import PoolEmbeddings as PoolEmbeddings
- Step 4: Run tests to verify they pass
Run: uv run pytest tests/llm/test_pool.py -v → PASS (8 tests). Ruff check + format.
- Step 5: Live smoke against the rack Ollama, then commit
uv run python -c "
from llm.pool import HostPool, embed_texts
p = HostPool(['http://127.0.0.1:11434'])
print(p.check('nomic-embed-text'))
print(len(embed_texts(p, 'nomic-embed-text', ['hello world'])[0]))
"
Expected: ['http://127.0.0.1:11434'] then 768.
git add src/llm tests/llm
git commit -m "feat(llm): multi-host Ollama pool — least-loaded embed fan-out (refs #568)"
Task 5: Chunker (#566)
Files:
- Create:
src/llm/chunk.py,tests/llm/test_chunk.py
Interfaces:
-
Produces (Tasks 6–7 consume):
@dataclass(frozen=True) Doc: key: str; text: str; metadata: dict[str, str]@dataclass(frozen=True) Chunk: id: str; text: str; metadata: dict[str, str]chunk_doc(doc: Doc, *, target_chars: int = 2000, overlap_chars: int = 200) -> list[Chunk]content_hash(text: str) -> str(sha256 hex, also the resume key inindex_state)- Chunk id format:
f"{doc.key}:{content_hash(doc.text)[:12]}:{seq:04d}"; chunk metadata = doc metadata +{"item_key": doc.key, "seq": str(seq)}.
-
Step 1: Write the failing tests
tests/llm/test_chunk.py:
"""llm.chunk — markdown-aware chunking with deterministic ids."""
from llm.chunk import Doc, chunk_doc, content_hash
FRONTMATTER_DOC = """---
comment_id: CMS-2019-0111-0042
docket_id: CMS-2019-0111
---
# Re: CY 2020 PFS Proposed Rule
We object to the E/M consolidation.
## Telehealth
Originating-site rules should be relaxed.
"""
def _doc(text, key="K1"):
return Doc(key=key, text=text, metadata={"docket": "CMS-2019-0111"})
class TestChunkDoc:
def test_strips_yaml_frontmatter(self):
chunks = chunk_doc(_doc(FRONTMATTER_DOC))
assert "comment_id:" not in chunks[0].text
assert chunks[0].text.startswith("# Re:")
def test_deterministic_ids(self):
a = chunk_doc(_doc(FRONTMATTER_DOC))
b = chunk_doc(_doc(FRONTMATTER_DOC))
assert [c.id for c in a] == [c.id for c in b]
h = content_hash(FRONTMATTER_DOC)
assert a[0].id == f"K1:{h[:12]}:0000"
def test_metadata_carries_item_key_and_seq(self):
chunks = chunk_doc(_doc(FRONTMATTER_DOC))
assert chunks[0].metadata["item_key"] == "K1"
assert chunks[0].metadata["docket"] == "CMS-2019-0111"
assert chunks[0].metadata["seq"] == "0"
def test_splits_on_headings_before_size(self):
text = "# A\n\n" + "para. " * 100 + "\n\n# B\n\nshort."
chunks = chunk_doc(_doc(text), target_chars=300)
assert all(len(c.text) <= 300 + 200 for c in chunks)
assert any(c.text.lstrip().startswith("# B") for c in chunks)
def test_long_paragraph_hard_wrapped_with_overlap(self):
text = "x" * 5000
chunks = chunk_doc(_doc(text), target_chars=2000, overlap_chars=200)
assert len(chunks) == 3
assert chunks[1].text[:200] == chunks[0].text[-200:]
def test_empty_and_whitespace_yield_nothing(self):
assert chunk_doc(_doc("")) == []
assert chunk_doc(_doc(" \n\n ")) == []
- Step 2: Run tests to verify they fail
Run: uv run pytest tests/llm/test_chunk.py -v
Expected: FAIL with ModuleNotFoundError: No module named 'llm.chunk'
- Step 3: Write the implementation
src/llm/chunk.py:
"""Markdown-aware chunking with deterministic ids.
Ids are ``{item_key}:{content_hash[:12]}:{seq:04d}`` — same input, same
ids — so pgvector upserts are idempotent and ``index_state`` can skip
unchanged items by comparing hashes.
Sizes are in characters (~4 chars/token; 2000 chars ≈ 500 tokens, inside
every candidate embed model's window).
"""
from __future__ import annotations
import hashlib
import re
from dataclasses import dataclass
@dataclass(frozen=True)
class Doc:
key: str
text: str
metadata: dict[str, str]
@dataclass(frozen=True)
class Chunk:
id: str
text: str
metadata: dict[str, str]
def content_hash(text: str) -> str:
return hashlib.sha256(text.encode()).hexdigest()
_FRONTMATTER = re.compile(r"\A---\n.*?\n---\n", re.DOTALL)
_HEADING = re.compile(r"^#{1,6} ", re.MULTILINE)
def _sections(text: str) -> list[str]:
"""Split at markdown headings, keeping the heading with its body."""
starts = [m.start() for m in _HEADING.finditer(text)]
if not starts:
return [text]
bounds = ([0] if starts[0] != 0 else []) + starts + [len(text)]
return [text[a:b] for a, b in zip(bounds, bounds[1:])]
def _pack(section: str, target: int, overlap: int) -> list[str]:
"""Greedily pack paragraphs; hard-wrap oversized ones with overlap."""
paras = [p for p in re.split(r"\n\n+", section) if p.strip()]
pieces: list[str] = []
buf = ""
for para in paras:
if buf and len(buf) + len(para) + 2 > target:
pieces.append(buf)
buf = ""
if len(para) > target:
step = target - overlap
for i in range(0, len(para), step):
pieces.append(para[max(0, i - overlap) if i else 0 : i + step + (overlap if i else 0)])
continue
buf = f"{buf}\n\n{para}" if buf else para
if buf:
pieces.append(buf)
return pieces
def chunk_doc(
doc: Doc, *, target_chars: int = 2000, overlap_chars: int = 200
) -> list[Chunk]:
body = _FRONTMATTER.sub("", doc.text).strip()
if not body:
return []
prefix = f"{doc.key}:{content_hash(doc.text)[:12]}"
texts = [
piece
for section in _sections(body)
for piece in _pack(section, target_chars, overlap_chars)
]
return [
Chunk(
id=f"{prefix}:{seq:04d}",
text=piece,
metadata={**doc.metadata, "item_key": doc.key, "seq": str(seq)},
)
for seq, piece in enumerate(texts)
]
Note: test_long_paragraph_hard_wrapped_with_overlap pins exact overlap
semantics — get the hard-wrap slice math right against the test (first
window [0 : step + 0]… adjust implementation until the three-chunk /
200-char-overlap assertions pass, keeping windows ≤ target+overlap).
- Step 4: Run tests to verify they pass
Run: uv run pytest tests/llm/test_chunk.py -v → PASS (6 tests). Ruff check + format.
- Step 5: Commit
git add src/llm/chunk.py tests/llm/test_chunk.py
git commit -m "feat(llm): markdown-aware chunker with deterministic ids (refs #566)"
Task 6: Document sources — comments + corpus (#566)
Files:
- Create:
src/llm/source.py,tests/llm/test_source.py
Interfaces:
-
Consumes:
llm.chunk.Doc;bib.store.Store(list_items(tag=..., limit=...),store._con()for attachments/URL SQL — the established convention, seesrc/cli/comments.py:75);rex.comments.combine.parse_combined; extraction layout.state/comments/<docket>/<comment_id>/{combined.md, <attachment>.md}. -
Produces (Task 7 consumes):
iter_comment_docs(store: Store, *, docket: str = "", root: Path | None = None) -> Iterator[Doc]— oneDocper comment; text = combined.md body when extracted, else bibabstract; metadata keys:docket,comment_id,doctype(="comment"),year(fromyear:tag when present).iter_corpus_docs(store: Store, *, tag: str = "") -> Iterator[Doc]— non-comment items; text = extracted attachment text (PyMuPDF viarex) joined with abstract; metadata:doctype(= item_type),year.comment_key_map(store: Store, docket: str) -> dict[str, tuple[str, str]]— comment_id → (bib item key, year) viaitems.url LIKE 'https://www.regulations.gov/comment/%'.
-
Step 1: Write the failing tests
tests/llm/test_source.py (in-memory bib Store + tmp extraction tree — no mocks needed for bib):
"""llm.source — comment + corpus Doc iterators."""
from pathlib import Path
import pytest
from bib.item import Item
from bib.store import Store
from llm.source import comment_key_map, iter_comment_docs, iter_corpus_docs
DOCKET = "CMS-2019-0111"
CID = f"{DOCKET}-0042"
COMBINED = f"""---
comment_id: {CID}
docket_id: {DOCKET}
---
We object to the E/M consolidation.
"""
@pytest.fixture
def store():
s = Store(":memory:")
key = s.create(
Item(
item_type="report",
title="A comment",
url=f"https://www.regulations.gov/comment/{CID}",
abstract="Inline abstract body.",
)
)
for tag in (f"docket:{DOCKET}", "doctype:comment", "year:2019"):
s.add_tag(key, tag)
s._comment_key = key # test convenience
return s
@pytest.fixture
def root(tmp_path):
d = tmp_path / DOCKET / CID
d.mkdir(parents=True)
(d / "combined.md").write_text(COMBINED)
return tmp_path
class TestCommentDocs:
def test_extracted_comment_uses_combined_body(self, store, root):
docs = list(iter_comment_docs(store, docket=DOCKET, root=root))
assert len(docs) == 1
assert docs[0].key == store._comment_key
assert "E/M consolidation" in docs[0].text
assert docs[0].metadata == {
"docket": DOCKET,
"comment_id": CID,
"doctype": "comment",
"year": "2019",
}
def test_unextracted_comment_falls_back_to_abstract(self, store, tmp_path):
docs = list(iter_comment_docs(store, docket=DOCKET, root=tmp_path))
assert docs[0].text == "Inline abstract body."
def test_docket_filter_excludes_others(self, store, root):
assert list(iter_comment_docs(store, docket="CMS-2021-0119", root=root)) == []
class TestCommentKeyMap:
def test_maps_comment_id_to_key_and_year(self, store):
mapping = comment_key_map(store, DOCKET)
assert mapping[CID] == (store._comment_key, "2019")
class TestCorpusDocs:
def test_non_comment_item_with_abstract(self, store):
key = store.create(
Item(item_type="rule", title="Final rule", abstract="Rule text.")
)
store.add_tag(key, "year:2020")
docs = list(iter_corpus_docs(store))
assert [d.key for d in docs] == [key]
assert docs[0].text == "Rule text."
assert docs[0].metadata["doctype"] == "rule"
def test_comments_excluded_from_corpus(self, store):
assert all(
d.metadata["doctype"] != "comment" for d in iter_corpus_docs(store)
)
Adjust the Item(...) constructor kwargs to the real bib.item.Item model
(read src/bib/item.py first — use whatever concrete subclass or field
names it defines; the assertions above are the contract, not the fixture
syntax).
- Step 2: Run tests to verify they fail
Run: uv run pytest tests/llm/test_source.py -v
Expected: FAIL with ModuleNotFoundError: No module named 'llm.source'
- Step 3: Write the implementation
src/llm/source.py:
"""Doc iterators over the bib store + comment extraction tree.
Comments: text prefers the #253 extraction output
(``.state/comments/<docket>/<comment_id>/combined.md`` body via
``rex.comments.combine.parse_combined``); falls back to the bib
``abstract`` for attachment-less comments. Corpus: every non-comment item;
attachment text extracted with the same rex extractors, joined with the
abstract.
"""
from __future__ import annotations
from pathlib import Path
from typing import Iterator
from bib.store import Store
from llm.chunk import Doc
_COMMENT_URL_PREFIX = "https://www.regulations.gov/comment/"
def _default_root() -> Path:
from conf import ROOT
return ROOT / ".state" / "comments"
def _year_of(store: Store, item_key: str) -> str:
row = store._con().execute(
"SELECT t.name FROM tags t "
"JOIN item_tags it ON it.tag_id = t.id "
"JOIN items i ON i.id = it.item_id "
"WHERE i.key = ? AND t.name LIKE 'year:%'",
(item_key,),
).fetchone()
return row[0].split(":", 1)[1] if row else ""
def comment_key_map(store: Store, docket: str) -> dict[str, tuple[str, str]]:
"""comment_id -> (bib item key, year) for one docket."""
rows = store._con().execute(
"SELECT i.key, i.url FROM items i "
"JOIN item_tags it ON it.item_id = i.id "
"JOIN tags t ON t.id = it.tag_id "
"WHERE t.name = ? AND i.url LIKE ?",
(f"docket:{docket}", _COMMENT_URL_PREFIX + "%"),
).fetchall()
return {
url.rsplit("/", 1)[-1]: (key, _year_of(store, key))
for key, url in rows
}
def iter_comment_docs(
store: Store, *, docket: str = "", root: Path | None = None
) -> Iterator[Doc]:
"""One Doc per comment: extraction body, else abstract fallback."""
from rex.comments.combine import parse_combined
root = root if root is not None else _default_root()
dockets = [docket] if docket else sorted(
t["name"].split(":", 1)[1]
for t in store.list_tags(namespace="docket")
if t["name"] != "docket:"
)
for dk in dockets:
for comment_id, (key, year) in sorted(comment_key_map(store, dk).items()):
meta = {
"docket": dk,
"comment_id": comment_id,
"doctype": "comment",
"year": year,
}
combined = root / dk / comment_id / "combined.md"
if combined.exists():
_, body = parse_combined(combined.read_text())
if body.strip():
yield Doc(key=key, text=body, metadata=meta)
continue
item = store.get(key)
if item.abstract.strip():
yield Doc(key=key, text=item.abstract, metadata=meta)
def _attachment_text(store: Store, item_key: str) -> str:
from rex.comments.combine import extract_attachment
rows = store._con().execute(
"SELECT a.storage_path FROM attachments a "
"JOIN items i ON i.id = a.item_id WHERE i.key = ?",
(item_key,),
).fetchall()
parts = []
for (storage_path,) in rows:
path = Path(storage_path)
if path.exists():
text, status = extract_attachment(path)
if status == "ok":
parts.append(text)
return "\n\n".join(parts)
def iter_corpus_docs(store: Store, *, tag: str = "") -> Iterator[Doc]:
"""Every non-comment item; attachment text + abstract."""
for item in store.list_items(tag=tag):
if "doctype:comment" in item.tags:
continue
text = "\n\n".join(
p for p in (_attachment_text(store, item.key), item.abstract) if p.strip()
)
if not text.strip():
continue
yield Doc(
key=item.key,
text=text,
metadata={
"doctype": item.item_type,
"year": _year_of(store, item.key),
},
)
Check against reality while implementing: (a) Item.tags — read
src/bib/item.py for how tags are exposed on a hydrated item; if
store.get() doesn't hydrate tags, use a _year_of-style SQL exclusion
(NOT EXISTS ... t.name = 'doctype:comment') instead. (b)
extract_attachment signature — read src/rex/comments/combine.py:135
(it may return (text, status) or a richer object; adapt the unpacking).
(c) list_items is O(n) get() calls — for the corpus (≈17k non-comment
items) this is acceptable; comments never go through it.
- Step 4: Run tests to verify they pass
Run: uv run pytest tests/llm/test_source.py -v → PASS. Ruff check + format.
- Step 5: Live sanity check on a real docket, then commit
uv run python -c "
from conf.connect import bib
from llm.source import iter_comment_docs
import itertools
docs = list(itertools.islice(iter_comment_docs(bib(), docket='CMS-2017-0092'), 5))
for d in docs: print(d.key, d.metadata['comment_id'], len(d.text))
"
Expected: 5 rows, nonzero lengths.
git add src/llm/source.py tests/llm/test_source.py
git commit -m "feat(llm): comment + corpus Doc sources with abstract fallback (refs #566)"
Task 7: Incremental indexer + stack llm index CLI (#567)
Files:
- Create:
src/llm/index.py,tests/llm/test_index.py - Create:
src/cli/llm.py - Modify:
src/cli/__init__.py(register subcommand)
Interfaces:
-
Consumes:
llm.config,llm.migrate.{migrate, ensure_hnsw},llm.pool.{HostPool, PoolEmbeddings, embed_texts},llm.chunk.{chunk_doc, content_hash},llm.source.{iter_comment_docs, iter_corpus_docs},langchain_postgres.PGVector. -
Produces:
vectorstore(collection: str, cfg: LlmConfig, pool: HostPool) -> PGVector— constructed withembedding_length=cfg.embed_dim,use_jsonb=True.index_docs(docs: Iterable[Doc], *, collection: str, cfg: LlmConfig, pool: HostPool, force: bool = False) -> dictreturning{"indexed": int, "skipped": int, "chunks": int}.- CLI:
stack llm index [--collection comments|corpus] [--docket ID] [--force] [--limit N].
-
Step 1: Write the failing tests
tests/llm/test_index.py — unit-level: mock PGVector, engine, and embed; assert skip/re-embed/delete-then-add logic:
"""llm.index — incremental, resumable embedding indexer."""
from unittest.mock import MagicMock, patch
from llm.chunk import Doc, content_hash
from llm.config import LlmConfig
from llm.index import index_docs
CFG = LlmConfig(
ollama_hosts=("http://h1:11434",),
embed_model="m",
instruct_model="g",
embed_dim=768,
pg_host="x",
pg_port=5432,
pg_db="llm",
pg_user="llm",
)
DOC = Doc(key="K1", text="Some body text.", metadata={"docket": "D"})
def _run(docs, state_rows, force=False):
"""Run index_docs with everything external mocked; return mocks."""
store = MagicMock()
engine = MagicMock()
conn = engine.begin.return_value.__enter__.return_value
conn.execute.return_value.fetchall.return_value = state_rows
with (
patch("llm.index._engine", return_value=engine),
patch("llm.index.vectorstore", return_value=store),
patch("llm.index.embed_texts", return_value=[[0.0] * 3]),
patch("llm.index.HostPool") as MockPool,
):
MockPool.return_value.check.return_value = ["http://h1:11434"]
stats = index_docs(
docs, collection="comments", cfg=CFG,
pool=MockPool.return_value, force=force,
)
return stats, store, conn
class TestIndexDocs:
def test_new_doc_embedded_and_recorded(self):
stats, store, conn = _run([DOC], state_rows=[])
assert stats == {"indexed": 1, "skipped": 0, "chunks": 1}
store.add_embeddings.assert_called_once()
kwargs = store.add_embeddings.call_args.kwargs
assert kwargs["ids"][0].startswith("K1:")
def test_unchanged_doc_skipped(self):
h = content_hash(DOC.text)
stats, store, _ = _run([DOC], state_rows=[("K1", h)])
assert stats["skipped"] == 1
store.add_embeddings.assert_not_called()
def test_changed_doc_deletes_old_chunks_first(self):
stats, store, conn = _run([DOC], state_rows=[("K1", "stalehash")])
assert stats["indexed"] == 1
deletes = " ".join(
str(c.args[0]) for c in conn.execute.call_args_list
)
assert "item_key" in deletes # DELETE ... cmetadata->>'item_key'
def test_force_reembeds_unchanged(self):
h = content_hash(DOC.text)
stats, store, _ = _run([DOC], state_rows=[("K1", h)], force=True)
assert stats["indexed"] == 1
def test_empty_doc_counts_skipped(self):
empty = Doc(key="K2", text=" ", metadata={})
stats, store, _ = _run([empty], state_rows=[])
assert stats == {"indexed": 0, "skipped": 1, "chunks": 0}
- Step 2: Run tests to verify they fail
Run: uv run pytest tests/llm/test_index.py -v
Expected: FAIL with ModuleNotFoundError: No module named 'llm.index'
- Step 3: Write the implementation
src/llm/index.py:
"""Incremental, resumable embedding indexer.
Per doc: compare ``content_hash`` against ``index_state``; skip when
unchanged, else delete the item's old chunks (by ``cmetadata->>
'item_key'``), embed via the host pool, upsert with deterministic ids,
and record the new hash — one transaction per doc, so a killed run
resumes exactly where it stopped.
"""
from __future__ import annotations
import logging
from typing import Iterable
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from llm.chunk import Doc, chunk_doc, content_hash
from llm.config import LlmConfig, pg_url
from llm.migrate import ensure_hnsw, migrate
from llm.pool import HostPool, PoolEmbeddings, embed_texts
log = logging.getLogger(__name__)
def _engine(cfg: LlmConfig) -> Engine:
return create_engine(pg_url(cfg))
def vectorstore(collection: str, cfg: LlmConfig, pool: HostPool):
from langchain_postgres import PGVector
return PGVector(
embeddings=PoolEmbeddings(pool, cfg.embed_model),
collection_name=collection,
connection=pg_url(cfg),
embedding_length=cfg.embed_dim,
use_jsonb=True,
)
def _state(engine: Engine, collection: str) -> dict[str, str]:
with engine.begin() as conn:
rows = conn.execute(
text(
"SELECT item_key, content_hash FROM index_state "
"WHERE collection = :c"
),
{"c": collection},
).fetchall()
return dict(rows)
def index_docs(
docs: Iterable[Doc],
*,
collection: str,
cfg: LlmConfig,
pool: HostPool,
force: bool = False,
) -> dict:
engine = _engine(cfg)
migrate(engine)
pool.check(cfg.embed_model)
store = vectorstore(collection, cfg, pool)
seen = _state(engine, collection)
stats = {"indexed": 0, "skipped": 0, "chunks": 0}
for doc in docs:
chunks = chunk_doc(doc)
if not chunks:
stats["skipped"] += 1
continue
h = content_hash(doc.text)
if not force and seen.get(doc.key) == h:
stats["skipped"] += 1
continue
vectors = embed_texts(
pool, cfg.embed_model, [c.text for c in chunks]
)
with engine.begin() as conn:
conn.execute(
text(
"DELETE FROM langchain_pg_embedding "
"WHERE cmetadata->>'item_key' = :k"
),
{"k": doc.key},
)
store.add_embeddings(
texts=[c.text for c in chunks],
embeddings=vectors,
metadatas=[c.metadata for c in chunks],
ids=[c.id for c in chunks],
)
with engine.begin() as conn:
conn.execute(
text(
"INSERT INTO index_state "
"(item_key, collection, content_hash, chunk_count) "
"VALUES (:k, :c, :h, :n) "
"ON CONFLICT (item_key, collection) DO UPDATE SET "
"content_hash = :h, chunk_count = :n, indexed_at = now()"
),
{"k": doc.key, "c": collection, "h": h, "n": len(chunks)},
)
stats["indexed"] += 1
stats["chunks"] += len(chunks)
if stats["indexed"] % 100 == 0:
log.info("indexed %(indexed)s (+%(chunks)s chunks)", stats)
ensure_hnsw(engine, cfg.embed_dim)
return stats
Implementation notes to verify against the installed langchain_postgres
version while coding: the delete-before-add is belt-and-braces (ids are
deterministic and add_embeddings upserts on id-conflict, but a changed
doc can shrink its chunk count, stranding stale ids — hence the metadata
delete). If PGVector in the pinned version lacks embedding_length,
drop the kwarg and rely on ensure_hnsw's ALTER ... TYPE vector(dim).
The first DELETE before any add_embeddings will also fail if
langchain_pg_embedding doesn't exist yet — create the store tables first
via PGVector.create_tables_if_not_exists() (or catch
ProgrammingError and skip); check which the pinned version offers.
src/cli/llm.py:
"""stack llm — local RAG over comments + corpus (P33: index only)."""
import typer
app = typer.Typer(no_args_is_help=True)
@app.command()
def index(
collection: str = typer.Option(
"comments", help="Which collection: comments | corpus."
),
docket: str = typer.Option("", help="Limit comments to one docket id."),
force: bool = typer.Option(False, help="Re-embed even when unchanged."),
limit: int = typer.Option(0, help="Stop after N docs (0 = all)."),
) -> None:
"""Embed comments/corpus into pgvector (incremental, resumable)."""
import itertools
from conf.connect import bib
from llm import config as llm_config
from llm.index import index_docs
from llm.pool import HostPool
from llm.source import iter_comment_docs, iter_corpus_docs
cfg = llm_config.load()
store = bib()
if collection == "comments":
docs = iter_comment_docs(store, docket=docket)
elif collection == "corpus":
docs = iter_corpus_docs(store)
else:
raise typer.BadParameter("collection must be 'comments' or 'corpus'")
if limit:
docs = itertools.islice(docs, limit)
stats = index_docs(
docs,
collection=collection,
cfg=cfg,
pool=HostPool.from_config(cfg),
force=force,
)
typer.echo(
f"indexed={stats['indexed']} skipped={stats['skipped']} "
f"chunks={stats['chunks']}"
)
src/cli/__init__.py — add with the other imports/registrations
(alphabetical with the existing add_typer calls):
from cli.llm import app as llm_app
app.add_typer(llm_app, name="llm", help="Local RAG — index, search, tag.")
- Step 4: Run tests to verify they pass
Run: uv run pytest tests/llm/ -v → PASS. Also uv run pytest tests/cli/ -q (the CLI registry usually has a help-text test — update it if it asserts the command list). Ruff check + format.
- Step 5: Commit
git add src/llm/index.py tests/llm/test_index.py src/cli/llm.py src/cli/__init__.py
git commit -m "feat(llm): incremental pgvector indexer + stack llm index CLI (refs #567)"
Task 8: Pilot-docket run + model bake-off + spec update (#564, #567)
Files:
- Create:
dev/scripts/llm_bakeoff.py - Modify:
docs/superpowers/specs/2026-07-16-llm-module-design.md(record results),stack.toml(if models change)
Interfaces:
-
Consumes: everything above, live services.
-
Produces: an indexed pilot docket in pgvector; a bake-off report + model decision recorded in the spec.
-
Step 1: Index the pilot docket (smallest first)
uv run stack llm index --collection comments --docket CMS-2017-0092
# kill it partway (Ctrl-C) once, rerun, confirm skipped > 0 and it completes:
uv run stack llm index --collection comments --docket CMS-2017-0092
Expected second run: indexed=<remainder> skipped=<already-done> … — resumability proven. Then verify retrieval directly:
PW=$(grep '^LLM_DB_PASSWORD=' .env | cut -d= -f2-)
LLM_DB_PASSWORD=$PW uv run python -c "
import llm
from llm.index import vectorstore
from llm.pool import HostPool
cfg = llm.load()
pool = HostPool.from_config(cfg); pool.check(cfg.embed_model)
vs = vectorstore('comments', cfg, pool)
for doc, score in vs.similarity_search_with_score('telehealth originating site', k=3):
print(round(score, 3), doc.metadata['item_key'], doc.page_content[:80])
"
Expected: 3 topically-relevant comment chunks.
- Step 2: Write the bake-off harness
dev/scripts/llm_bakeoff.py — self-contained, writes a markdown report to stdout. Retrieval proxy task: for N sampled extracted comments, query = the comment's bib abstract, relevant = its own chunks; score recall@5 + MRR per embed model, plus embed throughput. Generation task: run a fixed closed-vocab tagging prompt over M sample comments per instruct model; print outputs + tokens/s for human review.
"""Bake off embed + instruct models for the llm module (#564).
Usage:
LLM_DB_PASSWORD=... uv run python dev/scripts/llm_bakeoff.py \
--docket CMS-2017-0092 --sample 100
Embed metric: for each sampled comment, embed its chunks under each
candidate model into an in-memory matrix; query with the comment's
abstract; recall@5 = fraction of comments whose own chunk ranks top-5,
plus MRR and chunks/sec. Generation: tag 15 sample comments with each
candidate against a 12-tag toy vocab; print outputs + tokens/s for
eyeball review. No pgvector writes — pure in-memory numpy.
"""
import argparse
import itertools
import json
import time
import httpx
import numpy as np
from conf.connect import bib
from llm import config as llm_config
from llm.chunk import chunk_doc
from llm.pool import HostPool, embed_texts
from llm.source import iter_comment_docs
EMBED_CANDIDATES = ["nomic-embed-text", "bge-m3"]
GEN_CANDIDATES = ["llama3.1:8b", "qwen2.5:7b-instruct", "qwen2.5:14b-instruct-q4_K_M"]
TOY_VOCAB = [
"telehealth", "e/m-coding", "payment-rates", "quality-measures",
"drug-pricing", "supervision", "rural-access", "documentation-burden",
"site-neutrality", "aco-policy", "opioids", "scope-of-practice",
]
PROMPT = (
"You are tagging a public comment on a CMS rule. Choose up to 3 tags "
"from exactly this list, JSON array only, no prose: {vocab}\n\n"
"COMMENT:\n{body}\n\nTAGS:"
)
def embed_bakeoff(pool, docs, sample):
picked = [d for d in itertools.islice(docs, sample) if d.metadata]
for model in EMBED_CANDIDATES:
pool.check(model)
chunks, owners, queries = [], [], []
store = bib()
for i, d in enumerate(picked):
cs = chunk_doc(d)[:4]
chunks += [c.text for c in cs]
owners += [i] * len(cs)
queries.append(store.get(d.key).abstract or d.text[:300])
t0 = time.time()
mat = np.array(embed_texts(pool, model, chunks))
rate = len(chunks) / (time.time() - t0)
qs = np.array(embed_texts(pool, model, queries))
mat = mat / np.linalg.norm(mat, axis=1, keepdims=True)
qs = qs / np.linalg.norm(qs, axis=1, keepdims=True)
owners_arr = np.array(owners)
hits, rr = 0, 0.0
for i, q in enumerate(qs):
order = np.argsort(mat @ q)[::-1]
ranks = np.where(owners_arr[order] == i)[0]
if len(ranks):
rr += 1 / (ranks[0] + 1)
hits += ranks[0] < 5
n = len(qs)
print(f"| {model} | {hits / n:.2%} | {rr / n:.3f} | {rate:.0f} c/s |")
def gen_bakeoff(pool, docs, hosts):
picked = list(itertools.islice(docs, 15))
for model in GEN_CANDIDATES:
try:
pool.check(model)
except RuntimeError as exc:
print(f"\n### {model}: SKIPPED — {exc}")
continue
print(f"\n### {model}")
for d in picked:
body = d.text[:4000]
t0 = time.time()
with pool.acquire() as host, httpx.Client(timeout=300) as c:
r = c.post(f"{host}/api/generate", json={
"model": model, "stream": False,
"prompt": PROMPT.format(vocab=json.dumps(TOY_VOCAB), body=body),
})
out = r.json()
tps = out.get("eval_count", 0) / max(out.get("eval_duration", 1) / 1e9, 0.001)
print(f"- {d.metadata['comment_id']} ({tps:.0f} tok/s): "
f"{out['response'].strip()[:120]}")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--docket", default="CMS-2017-0092")
ap.add_argument("--sample", type=int, default=100)
args = ap.parse_args()
cfg = llm_config.load()
pool = HostPool.from_config(cfg)
print("| model | recall@5 | MRR | throughput |")
print("|---|---|---|---|")
embed_bakeoff(pool, iter_comment_docs(bib(), docket=args.docket), args.sample)
gen_bakeoff(pool, iter_comment_docs(bib(), docket=args.docket), cfg.ollama_hosts)
if __name__ == "__main__":
main()
(dev/scripts are exempt from the coverage bar — confirm with grep -n "source" pyproject.toml: coverage source is src/ only.)
- Step 3: Run the bake-off
docker exec ollama ollama pull bge-m3
docker exec ollama ollama pull qwen2.5:7b-instruct
docker exec ollama ollama pull qwen2.5:14b-instruct-q4_K_M # ~9 GB — fits the 3060, tight
uv run python dev/scripts/llm_bakeoff.py --docket CMS-2017-0092 --sample 100 | tee /tmp/bakeoff.md
Review the table + generation outputs. Decide embed model (recall/MRR, dims, speed) and instruct model (tag validity — outputs parse as JSON arrays drawn from the vocab — and tok/s).
-
Step 4: Record the decision
-
If the winner differs from the defaults, update
stack.toml[llm]embed_model/embed_dim/instruct_modeland re-runuv run stack llm index --collection comments --docket CMS-2017-0092 --force(changed model ⇒ changed vectors). -
Append a
## Model bake-off (P33 #564)section todocs/superpowers/specs/2026-07-16-llm-module-design.mdwith the results table, generation notes, per-host sizing (3060 vs 4090/5080), and the decision. -
Step 5: Full-suite verification, commit, close out
uv run pytest tests/ -q -n auto # entire suite green
uv run ruff check src/ tests/ dev/scripts/llm_bakeoff.py
git add dev/scripts/llm_bakeoff.py docs/superpowers/specs/2026-07-16-llm-module-design.md stack.toml
git commit -m "feat(llm): model bake-off + pilot docket indexed; record decision (refs #564 #567)"
Then comment results on #564 and close #563–#568 as their deliverables land (issues auto-close only via closes #N in merged commit messages; otherwise close via the tracker with a comment linking the commit).
Self-Review Notes
- Spec coverage: #563 → Task 1; #565 → Task 2; #564 → Tasks 3 + 8; #566 → Tasks 5–6; #568 → Task 4; #567 → Tasks 7 + 8 (pilot). P33's "v1 gate: one docket end-to-end before corpus-wide run" → Task 8 Step 1.
- Known verify-at-execution points (flagged inline, not placeholders):
bib.item.Itemconstructor shape (Task 6),extract_attachmentreturn shape (Task 6),langchain_postgresembedding_length/create_tables_if_not_existsavailability (Task 7),api.auth bootstrap --onlysupport (Task 2). - Type consistency:
LlmConfigfields,Doc/Chunk,HostPool.acquire/check,embed_texts,index_docssignatures are used identically across Tasks 4–8.