homelab/stack
1
0
Fork 0
Code Issues 20 Pull Requests Actions Packages Projects 1 Releases Wiki Activity
Files
87cb0370d71588633e6c1a356b47836cdbd4295a
stack/docs/superpowers/plans/2026-07-16-llm-foundation-p33.md

1751 lines
60 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# 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=99` in CI); every new source file ships with tests in the same commit.
- Ruff: `select = ["E", "F", "I"]`, line-length 88. Run `uv 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` + `monkeypatch` only — **no respx, no pytest-mock** (not in dev deps).
- Tests live in `tests/llm/` mirroring `src/llm/`, with empty `tests/llm/__init__.py`.
- New module must be added to `[tool.uv.build-backend] module-name` in pyproject.toml or it won't package.
- `.gitea/workflows/*.yml` are generated — never hand-edit; regen via `uv run python dev/scripts/gen_config.py` if 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 at `127.0.0.1:5432`. Secrets come from `.env` via `os.environ`, never stack.toml.
- Ollama model names used throughout until the bake-off (Task 8) says otherwise: embed `nomic-embed-text` (768-dim), instruct `llama3.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) and `llm.config.load() -> LlmConfig`; `llm.config.pg_url(cfg: LlmConfig) -> str` returning a SQLAlchemy/psycopg URL. All later tasks import these exact names.
- [ ] **Step 1: Add config + packaging entries**
`stack.toml` — add after the `[bcda]` section:
```toml
[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:
```toml
ollama = "http://ollama:11434"
```
`pyproject.toml` — add to `[project.optional-dependencies]` (after `bls`):
```toml
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`:
```python
"""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):
```python
"""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`:
```python
"""[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`
```bash
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` (postgres `ports:`)
- Create: `src/llm/migrate.py`, `tests/llm/test_migrate.py`
- Modify: `tests/api/test_manifest.py` if 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; creates `index_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):
```python
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:`):
```yaml
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):
```python
"""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`:
```python
"""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**
```bash
# 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**
```bash
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`.
```bash
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` (new `ollama` service + `ollama_models` volume)
**Interfaces:**
- Produces: Ollama API at `http://127.0.0.1:11434` (host/LAN) and `http://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`):
```yaml
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**
```bash
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**
```bash
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-export `HostPool`)
**Interfaces:**
- Consumes: `llm.config.LlmConfig`.
- Produces (Tasks 7–8 and P34/P35 consume these exact names):
- `class HostPool` with `HostPool(hosts: Sequence[str])`, `HostPool.from_config(cfg: LlmConfig) -> HostPool`, `check(model: str) -> list[str]` (drops hosts that are down or lack the model; raises `RuntimeError` if 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`:
```python
"""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`:
```python
"""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`:
```python
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**
```bash
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`.
```bash
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 in `index_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`:
```python
"""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`:
```python
"""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**
```bash
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, see `src/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]` — one `Doc` per comment; text = combined.md body when extracted, else bib `abstract`; metadata keys: `docket`, `comment_id`, `doctype` (= `"comment"`), `year` (from `year:` tag when present).
- `iter_corpus_docs(store: Store, *, tag: str = "") -> Iterator[Doc]` — non-comment items; text = extracted attachment text (PyMuPDF via `rex`) 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) via `items.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):
```python
"""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`:
```python
"""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**
```bash
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.
```bash
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 with `embedding_length=cfg.embed_dim`, `use_jsonb=True`.
- `index_docs(docs: Iterable[Doc], *, collection: str, cfg: LlmConfig, pool: HostPool, force: bool = False) -> dict` returning `{"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:
```python
"""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`:
```python
"""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`:
```python
"""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):
```python
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**
```bash
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)**
```bash
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:
```bash
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.
```python
"""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**
```bash
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_model` and re-run `uv run stack llm index --collection comments --docket CMS-2017-0092 --force` (changed model ⇒ changed vectors).
- Append a `## Model bake-off (P33 #564)` section to `docs/superpowers/specs/2026-07-16-llm-module-design.md` with the results table, generation notes, per-host sizing (3060 vs 4090/5080), and the decision.
- [ ] **Step 5: Full-suite verification, commit, close out**
```bash
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.Item` constructor shape (Task 6), `extract_attachment` return shape (Task 6), `langchain_postgres` `embedding_length`/`create_tables_if_not_exists` availability (Task 7), `api.auth bootstrap --only` support (Task 2).
- **Type consistency:** `LlmConfig` fields, `Doc`/`Chunk`, `HostPool.acquire`/`check`, `embed_texts`, `index_docs` signatures are used identically across Tasks 4–8.
Reference in New Issue View Git Blame Copy Permalink