Files
stack/docs/superpowers/plans/2026-07-16-llm-foundation-p33.md

60 KiB
Raw Permalink Blame History

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:

[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 (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):

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 (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):

  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-export HostPool)

Interfaces:

  • Consumes: llm.config.LlmConfig.

  • Produces (Tasks 78 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:

"""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 67 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:

"""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, 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):

"""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 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:

"""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_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

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 56; #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 48.