feat(llm): theme tagging chain + docket runner + write-back — llm.tagger, Store.merge_extra, stack llm tag (refs #575, #576)

llm.tagger: cosine shortlist of the vocabulary's theme cards against a
comment's already-indexed chunks (no re-embedding — text and vector come
from langchain_pg_embedding), one yes/no judgement per candidate on the
largest live host at temperature 0 with the scoring chunk as evidence,
up to max_tags accepted with the shortlist similarity as confidence.
State on the bib item (extra_json.llm_tags: vocab version, model,
content hash, tags with evidence chunk, the full shortlist with
verdicts) makes runs resumable; apply_tags replaces exactly the item's
llm:* tags and never touches hand tags. Store.merge_extra merges
top-level keys into extra_json. Every collaborator is injected so the
pure pieces are tested without Ollama or pgvector.

stack llm tag --docket D [--limit N] [--force] [--dry-run] [--top]
[--max-tags] [--min-confidence] [--verbose]; _tagging_runtime wires the
store, pgvector loader, card vectors, judge and model.
This commit is contained in:
kert
2026-09-22 16:53:23 -04:00
parent 7693e97d36
commit 65f8d785a1
60 changed files with 1078 additions and 53 deletions

View File

@@ -1,6 +1,6 @@
--- ---
title: stack api serve title: stack api serve
sidebar_position: 62 sidebar_position: 63
--- ---
# `stack api serve` # `stack api serve`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack api title: stack api
sidebar_position: 61 sidebar_position: 62
--- ---
# `stack api` # `stack api`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack db comment title: stack db comment
sidebar_position: 55 sidebar_position: 56
--- ---
# `stack db comment` # `stack db comment`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack db inspect title: stack db inspect
sidebar_position: 56 sidebar_position: 57
--- ---
# `stack db inspect` # `stack db inspect`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack db title: stack db
sidebar_position: 54 sidebar_position: 55
--- ---
# `stack db` # `stack db`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack docs build title: stack docs build
sidebar_position: 58 sidebar_position: 59
--- ---
# `stack docs build` # `stack docs build`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack docs generate title: stack docs generate
sidebar_position: 60 sidebar_position: 61
--- ---
# `stack docs generate` # `stack docs generate`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack docs serve title: stack docs serve
sidebar_position: 59 sidebar_position: 60
--- ---
# `stack docs serve` # `stack docs serve`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack docs title: stack docs
sidebar_position: 57 sidebar_position: 58
--- ---
# `stack docs` # `stack docs`

38
docs/docs/cli/llm-tag.md Normal file
View File

@@ -0,0 +1,38 @@
---
title: stack llm tag
sidebar_position: 54
---
# `stack llm tag`
```
Usage: stack llm tag [OPTIONS]
Closed-vocabulary theme tagging of one docket's comments (P35): shortlist by
similarity to the theme cards, judge each candidate yes/no on the largest live
host with the scoring chunk as evidence, record state in the item's
extra_json, and replace its llm: tags. Resumable — unchanged comments are
skipped unless --force.
╭─ Options ────────────────────────────────────────────────────────────────────╮
│ * --docket TEXT regulations.gov docket id, e.g. │
│ CMS-2026-2377. │
│ [required] │
│ --limit INTEGER Stop after N comments (newest first; 0 = │
│ all). │
│ [default: 0] │
│ --force Re-tag comments whose stored state is │
│ current. │
│ --dry-run Compute and print, write nothing. │
│ --top INTEGER Themes shortlisted per comment before │
│ judging. │
│ [default: 8] │
│ --max-tags INTEGER Most themes written per comment. │
│ [default: 4] │
│ --min-confidence FLOAT Shortlist similarity below which a │
│ judged yes is not written. │
│ [default: 0.0] │
│ --verbose Print every comment's tags. │
│ --help Show this message and exit. │
╰──────────────────────────────────────────────────────────────────────────────╯
```

View File

@@ -23,5 +23,10 @@ Usage: stack llm [OPTIONS] COMMAND [ARGS]...
│ and list its slugs, or show one theme's definition, synonyms and │ │ and list its slugs, or show one theme's definition, synonyms and │
│ the FR │ │ the FR │
│ section stems it was seeded from. │ │ section stems it was seeded from. │
│ tag Closed-vocabulary theme tagging of one docket's comments (P35): │
│ shortlist by similarity to the theme cards, judge each candidate │
│ yes/no on the largest live host with the scoring chunk as evidence, │
│ record state in the item's extra_json, and replace its llm: tags. │
│ Resumable — unchanged comments are skipped unless --force. │
╰──────────────────────────────────────────────────────────────────────────────╯ ╰──────────────────────────────────────────────────────────────────────────────╯
``` ```

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail attach-smarthost title: stack mail attach-smarthost
sidebar_position: 103 sidebar_position: 104
--- ---
# `stack mail attach-smarthost` # `stack mail attach-smarthost`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail dkim-export title: stack mail dkim-export
sidebar_position: 102 sidebar_position: 103
--- ---
# `stack mail dkim-export` # `stack mail dkim-export`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail dns title: stack mail dns
sidebar_position: 101 sidebar_position: 102
--- ---
# `stack mail dns` # `stack mail dns`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail down title: stack mail down
sidebar_position: 99 sidebar_position: 100
--- ---
# `stack mail down` # `stack mail down`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail provision title: stack mail provision
sidebar_position: 97 sidebar_position: 98
--- ---
# `stack mail provision` # `stack mail provision`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail rotate-creds title: stack mail rotate-creds
sidebar_position: 104 sidebar_position: 105
--- ---
# `stack mail rotate-creds` # `stack mail rotate-creds`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail seed-mailboxes title: stack mail seed-mailboxes
sidebar_position: 105 sidebar_position: 106
--- ---
# `stack mail seed-mailboxes` # `stack mail seed-mailboxes`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail status title: stack mail status
sidebar_position: 100 sidebar_position: 101
--- ---
# `stack mail status` # `stack mail status`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail up title: stack mail up
sidebar_position: 98 sidebar_position: 99
--- ---
# `stack mail up` # `stack mail up`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail wire-git title: stack mail wire-git
sidebar_position: 106 sidebar_position: 107
--- ---
# `stack mail wire-git` # `stack mail wire-git`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack mail title: stack mail
sidebar_position: 96 sidebar_position: 97
--- ---
# `stack mail` # `stack mail`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack perf show title: stack perf show
sidebar_position: 64 sidebar_position: 65
--- ---
# `stack perf show` # `stack perf show`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack perf title: stack perf
sidebar_position: 63 sidebar_position: 64
--- ---
# `stack perf` # `stack perf`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs cpt-ingest title: stack pfs cpt-ingest
sidebar_position: 78 sidebar_position: 79
--- ---
# `stack pfs cpt-ingest` # `stack pfs cpt-ingest`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs elements title: stack pfs elements
sidebar_position: 70 sidebar_position: 71
--- ---
# `stack pfs elements` # `stack pfs elements`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs exposure title: stack pfs exposure
sidebar_position: 75 sidebar_position: 76
--- ---
# `stack pfs exposure` # `stack pfs exposure`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs families title: stack pfs families
sidebar_position: 72 sidebar_position: 73
--- ---
# `stack pfs families` # `stack pfs families`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs guidance title: stack pfs guidance
sidebar_position: 73 sidebar_position: 74
--- ---
# `stack pfs guidance` # `stack pfs guidance`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs lineage title: stack pfs lineage
sidebar_position: 71 sidebar_position: 72
--- ---
# `stack pfs lineage` # `stack pfs lineage`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs reaction title: stack pfs reaction
sidebar_position: 74 sidebar_position: 75
--- ---
# `stack pfs reaction` # `stack pfs reaction`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs review title: stack pfs review
sidebar_position: 77 sidebar_position: 78
--- ---
# `stack pfs review` # `stack pfs review`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs utilization title: stack pfs utilization
sidebar_position: 76 sidebar_position: 77
--- ---
# `stack pfs utilization` # `stack pfs utilization`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack pfs title: stack pfs
sidebar_position: 69 sidebar_position: 70
--- ---
# `stack pfs` # `stack pfs`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma eligible title: stack prisma eligible
sidebar_position: 90 sidebar_position: 91
--- ---
# `stack prisma eligible` # `stack prisma eligible`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma export title: stack prisma export
sidebar_position: 87 sidebar_position: 88
--- ---
# `stack prisma export` # `stack prisma export`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma extract title: stack prisma extract
sidebar_position: 91 sidebar_position: 92
--- ---
# `stack prisma extract` # `stack prisma extract`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma fetch title: stack prisma fetch
sidebar_position: 93 sidebar_position: 94
--- ---
# `stack prisma fetch` # `stack prisma fetch`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma flow title: stack prisma flow
sidebar_position: 92 sidebar_position: 93
--- ---
# `stack prisma flow` # `stack prisma flow`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma init title: stack prisma init
sidebar_position: 86 sidebar_position: 87
--- ---
# `stack prisma init` # `stack prisma init`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma ping-llm title: stack prisma ping-llm
sidebar_position: 88 sidebar_position: 89
--- ---
# `stack prisma ping-llm` # `stack prisma ping-llm`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma run title: stack prisma run
sidebar_position: 94 sidebar_position: 95
--- ---
# `stack prisma run` # `stack prisma run`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma screen title: stack prisma screen
sidebar_position: 89 sidebar_position: 90
--- ---
# `stack prisma screen` # `stack prisma screen`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma vpn title: stack prisma vpn
sidebar_position: 95 sidebar_position: 96
--- ---
# `stack prisma vpn` # `stack prisma vpn`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack prisma title: stack prisma
sidebar_position: 85 sidebar_position: 86
--- ---
# `stack prisma` # `stack prisma`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack rec list title: stack rec list
sidebar_position: 66 sidebar_position: 67
--- ---
# `stack rec list` # `stack rec list`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack rec opps title: stack rec opps
sidebar_position: 68 sidebar_position: 69
--- ---
# `stack rec opps` # `stack rec opps`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack rec pfs title: stack rec pfs
sidebar_position: 67 sidebar_position: 68
--- ---
# `stack rec pfs` # `stack rec pfs`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack rec title: stack rec
sidebar_position: 65 sidebar_position: 66
--- ---
# `stack rec` # `stack rec`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack zot dump-schema title: stack zot dump-schema
sidebar_position: 80 sidebar_position: 81
--- ---
# `stack zot dump-schema` # `stack zot dump-schema`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack zot fix-dates title: stack zot fix-dates
sidebar_position: 81 sidebar_position: 82
--- ---
# `stack zot fix-dates` # `stack zot fix-dates`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack zot fix-fields title: stack zot fix-fields
sidebar_position: 83 sidebar_position: 84
--- ---
# `stack zot fix-fields` # `stack zot fix-fields`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack zot fix-keys title: stack zot fix-keys
sidebar_position: 82 sidebar_position: 83
--- ---
# `stack zot fix-keys` # `stack zot fix-keys`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack zot verify-parity title: stack zot verify-parity
sidebar_position: 84 sidebar_position: 85
--- ---
# `stack zot verify-parity` # `stack zot verify-parity`

View File

@@ -1,6 +1,6 @@
--- ---
title: stack zot title: stack zot
sidebar_position: 79 sidebar_position: 80
--- ---
# `stack zot` # `stack zot`

View File

@@ -613,6 +613,24 @@ class Store:
con.commit() con.commit()
return added return added
def merge_extra(self, item_key: str, updates: dict[str, Any]) -> None:
"""Merge *updates* into the item's ``extra_json`` (top-level keys
replace; everything else is kept). A missing item is a no-op."""
con = self._con()
row = con.execute(
"SELECT extra_json FROM items WHERE key = ?", (item_key,)
).fetchone()
if row is None:
return
current = json.loads(row[0] or "{}") if row[0] else {}
current.update(updates)
con.execute(
"UPDATE items SET extra_json = ?, "
"updated_at = strftime('%Y-%m-%dT%H:%M:%SZ','now') WHERE key = ?",
(json.dumps(current), item_key),
)
con.commit()
def remove_tag(self, item_key: str, tag: str) -> None: def remove_tag(self, item_key: str, tag: str) -> None:
"""Remove a tag from an item.""" """Remove a tag from an item."""
con = self._con() con = self._con()

View File

@@ -1,5 +1,7 @@
"""stack llm — local RAG over the library (index + fleet + chat serve).""" """stack llm — local RAG over the library (index + fleet + chat serve)."""
from typing import Any
import typer import typer
app = typer.Typer(no_args_is_help=True) app = typer.Typer(no_args_is_help=True)
@@ -236,3 +238,103 @@ def vocab(
typer.echo( typer.echo(
f"vocabulary v{v.version}: {len(v)} themes, {len(v.retired)} retired ({v.source})" f"vocabulary v{v.version}: {len(v)} themes, {len(v.retired)} retired ({v.source})"
) )
def _tagging_runtime(cfg: Any) -> dict[str, Any]:
"""Everything `stack llm tag` needs from the live stack — one place
the tests monkeypatch: the bib store, the pgvector chunk loader, the
theme-card vectors (embedded once per run through the pool), the
yes/no judge on the largest live host, and the model name recorded
in each item's state."""
from conf.connect import bib
from llm.index import _engine
from llm.pool import HostPool, PoolEmbeddings, pick_model
from llm.tagger import card_vectors, make_judge, pg_chunks
from llm.vocab import load as load_vocab
vocab = load_vocab()
pool = HostPool.from_config(cfg)
pool.check(cfg.embed_model)
cards = card_vectors(vocab, PoolEmbeddings(pool, cfg.embed_model).embed_documents)
judge = make_judge(cfg, pool)
with pool.acquire_generation() as host:
model = pick_model(cfg, pool, host)
return {
"store": bib(),
"vocab": vocab,
"load_chunks": pg_chunks(_engine(cfg)),
"card_vecs": cards,
"judge": judge,
"model": model,
}
@app.command()
def tag(
docket: str = typer.Option(
..., "--docket", help="regulations.gov docket id, e.g. CMS-2026-2377."
),
limit: int = typer.Option(
0, "--limit", help="Stop after N comments (newest first; 0 = all)."
),
force: bool = typer.Option(
False, "--force", help="Re-tag comments whose stored state is current."
),
dry_run: bool = typer.Option(
False, "--dry-run", help="Compute and print, write nothing."
),
top: int = typer.Option(
8, "--top", help="Themes shortlisted per comment before judging."
),
max_tags: int = typer.Option(
4, "--max-tags", help="Most themes written per comment."
),
min_confidence: float = typer.Option(
0.0,
"--min-confidence",
help="Shortlist similarity below which a judged yes is not written.",
),
verbose: bool = typer.Option(
False, "--verbose", help="Print every comment's tags."
),
) -> None:
"""Closed-vocabulary theme tagging of one docket's comments (P35):
shortlist by similarity to the theme cards, judge each candidate
yes/no on the largest live host with the scoring chunk as evidence,
record state in the item's extra_json, and replace its llm: tags.
Resumable — unchanged comments are skipped unless --force."""
from llm import config as llm_config
from llm.tagger import run as run_tagging
cfg = llm_config.load()
rt = _tagging_runtime(cfg)
def progress(key: str, result: Any, why: str) -> None:
if result is None:
if verbose:
typer.echo(f"{key} {why}")
return
if verbose or dry_run:
typer.echo(f"{key} {', '.join(result.slugs) or '(none)'}")
stats = run_tagging(
rt["store"],
vocab=rt["vocab"],
docket=docket,
load_chunks=rt["load_chunks"],
card_vecs=rt["card_vecs"],
judge=rt["judge"],
model=rt["model"],
limit=limit,
force=force,
dry_run=dry_run,
top=top,
max_tags=max_tags,
min_confidence=min_confidence,
progress=progress,
)
typer.echo(
f"tag {docket} (vocab v{rt['vocab'].version}, {rt['model']}{', dry run' if dry_run else ''}): "
f"seen={stats['seen']} tagged={stats['tagged']} unchanged={stats['skipped_state']} "
f"no_chunks={stats['no_chunks']} tags_written={stats['tags_written']}"
)

369
src/llm/tagger.py Normal file
View File

@@ -0,0 +1,369 @@
"""Closed-vocabulary theme tagging of rulemaking comments (P35, #575/#576).
For one comment the chain is: its already-indexed chunks (text + vector,
``langchain_pg_embedding`` ``comments`` collection) → cosine similarity
against every theme card (``llm.vocab.Theme.card`` embedded once per run)
→ the ``top`` themes shortlisted, each with the chunk that scored it →
one yes/no judgement per candidate on the largest live host, with that
chunk as the evidence → up to ``max_tags`` accepted themes, confidence =
the shortlist similarity. Nothing outside ``{yes, no}`` counts as ``no``.
State lives on the bib item (``extra_json["llm_tags"]``: vocab version,
model, content hash, the accepted tags with confidence and evidence
chunk, and the shortlist with every verdict) so a run is resumable and a
re-run skips comments whose version + model + content are unchanged.
Write-back replaces exactly the item's ``llm:*`` tags and never touches
hand tags.
Every collaborator is injected (``embed``, ``judge``, ``load_chunks``) so
the pure pieces run in tests without Ollama or pgvector; the CLI wires
the real ones (``make_judge``, ``card_vectors``, ``pg_chunks``).
"""
from __future__ import annotations
import hashlib
import json
import logging
import math
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from typing import Any, Callable, Iterable, Mapping, Sequence
from llm.vocab import Theme, Vocab, split_tags
log = logging.getLogger(__name__)
Vector = Sequence[float]
Embed = Callable[[Sequence[str]], list[list[float]]]
Judge = Callable[[Theme, str], bool | None]
ChunkLoader = Callable[[str], list[tuple[int, str, list[float]]]]
STATE_KEY = "llm_tags"
@dataclass(frozen=True)
class Candidate:
slug: str
score: float
chunk_idx: int
@dataclass(frozen=True)
class TagInfo:
confidence: float
evidence_chunk: int
@dataclass(frozen=True)
class TagResult:
key: str
tags: dict[str, TagInfo]
shortlisted: tuple[Candidate, ...]
judged: dict[str, bool | None]
content_hash: str
@property
def slugs(self) -> list[str]:
return sorted(self.tags, key=lambda s: -self.tags[s].confidence)
# ── pure pieces ───────────────────────────────────────────────────────
def cosine(a: Vector, b: Vector) -> float:
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
return dot / (na * nb) if na and nb else 0.0
def shortlist(
chunk_vecs: Sequence[Vector], card_vecs: Mapping[str, Vector], *, top: int = 8
) -> list[Candidate]:
"""The ``top`` themes by their best cosine similarity to any chunk,
each with the index of the chunk that scored it (the evidence)."""
if not chunk_vecs:
return []
out: list[Candidate] = []
for slug, cv in card_vecs.items():
best_i, best = 0, -1.0
for i, v in enumerate(chunk_vecs):
s = cosine(v, cv)
if s > best:
best_i, best = i, s
out.append(Candidate(slug, round(best, 4), best_i))
out.sort(key=lambda c: (-c.score, c.slug))
return out[:top]
_SYSTEM = (
"You decide whether a public comment letter on a Medicare physician fee "
"schedule rule addresses a given theme. Answer with exactly one word: yes "
"or no. Say yes only when the excerpt actually discusses the theme, not "
"when it merely mentions a related word in passing."
)
def judge_prompt(theme: Theme, evidence: str) -> list[dict]:
return [
{"role": "system", "content": _SYSTEM},
{
"role": "user",
"content": (
f"Theme: {theme.label}\nDefinition: {theme.definition}\n"
f"Phrases letters use: {', '.join(theme.synonyms) or '—'}\n\n"
f"Excerpt from the comment:\n{evidence.strip()}\n\n"
"Does this comment address the theme? Answer yes or no."
),
},
]
def parse_yes_no(reply: str) -> bool | None:
word = (reply or "").strip().strip(".!\"'").lower().split()
if not word:
return None
if word[0] in ("yes", "y"):
return True
if word[0] in ("no", "n"):
return False
return None
def content_hash(chunks: Sequence[str]) -> str:
h = hashlib.sha1() # noqa: S324 — change detection, not security
for c in chunks:
h.update(c.encode("utf-8", "replace"))
h.update(b"\x00")
return h.hexdigest()[:16]
def tag_comment(
key: str,
chunks: Sequence[str],
chunk_vecs: Sequence[Vector],
*,
vocab: Vocab,
card_vecs: Mapping[str, Vector],
judge: Judge,
top: int = 8,
max_tags: int = 4,
min_confidence: float = 0.0,
) -> TagResult:
cands = shortlist(chunk_vecs, card_vecs, top=top)
judged: dict[str, bool | None] = {}
tags: dict[str, TagInfo] = {}
for c in cands:
if c.slug not in vocab:
continue
verdict = judge(vocab.get(c.slug), chunks[c.chunk_idx])
judged[c.slug] = verdict
if verdict and c.score >= min_confidence and len(tags) < max_tags:
tags[c.slug] = TagInfo(confidence=c.score, evidence_chunk=c.chunk_idx)
return TagResult(key, tags, tuple(cands), judged, content_hash(chunks))
def state_matches(
extra: Mapping[str, Any], *, version: int, model: str, digest: str
) -> bool:
st = extra.get(STATE_KEY) or {}
return (
st.get("version") == version
and st.get("model") == model
and st.get("content_hash") == digest
)
def state_payload(result: TagResult, *, version: int, model: str) -> dict:
return {
"version": version,
"model": model,
"content_hash": result.content_hash,
"tagged_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
"tags": {s: asdict(i) for s, i in result.tags.items()},
"shortlist": [
{
"slug": c.slug,
"score": c.score,
"chunk": c.chunk_idx,
"verdict": result.judged.get(c.slug),
}
for c in result.shortlisted
],
}
# ── bib side ──────────────────────────────────────────────────────────
def apply_tags(
store: Any, key: str, slugs: Iterable[str], *, vocab: Vocab
) -> tuple[int, int]:
"""Make the item's ``llm:*`` tags exactly ``slugs``; hand tags untouched.
Returns ``(added, removed)``."""
want = {vocab.tag(s) for s in slugs}
have, _hand = split_tags(store.get(key).tags)
removed = 0
for t in have:
if t not in want:
store.remove_tag(key, t)
removed += 1
added = store.add_tags(key, sorted(want - set(have))) if want - set(have) else 0
return added, removed
def _extra(store: Any, key: str) -> dict:
row = (
store._con()
.execute( # noqa: SLF001
"SELECT extra_json FROM items WHERE key = ?", (key,)
)
.fetchone()
)
return json.loads(row[0] or "{}") if row and row[0] else {}
def docket_keys(store: Any, docket: str, *, limit: int = 0) -> list[str]:
"""Comment item keys in *docket*, newest posted first."""
con = store._con() # noqa: SLF001
sql = (
"SELECT i.key FROM items i "
"WHERE i.id IN (SELECT item_id FROM item_tags WHERE tag_id IN "
"(SELECT id FROM tags WHERE name = ?)) "
"AND i.url LIKE 'https://www.regulations.gov/comment/%' "
"ORDER BY i.date_published DESC, i.id DESC"
)
if limit:
sql += f" LIMIT {int(limit)}"
return [r[0] for r in con.execute(sql, (f"reg-docket:{docket}",)).fetchall()]
def run(
store: Any,
*,
vocab: Vocab,
docket: str,
load_chunks: ChunkLoader,
card_vecs: Mapping[str, Vector],
judge: Judge,
model: str,
limit: int = 0,
force: bool = False,
dry_run: bool = False,
top: int = 8,
max_tags: int = 4,
min_confidence: float = 0.0,
progress: Callable[[str, TagResult | None, str], None] | None = None,
) -> dict[str, int]:
"""Tag every comment in *docket* (newest first); resumable via the
item's stored state; ``dry_run`` computes and reports but writes
nothing."""
stats = {
"seen": 0,
"tagged": 0,
"skipped_state": 0,
"no_chunks": 0,
"tags_written": 0,
}
for key in docket_keys(store, docket, limit=limit):
stats["seen"] += 1
extra = _extra(store, key)
rows = load_chunks(key)
if not rows:
stats["no_chunks"] += 1
if progress:
progress(key, None, "no-chunks")
continue
rows = sorted(rows, key=lambda r: r[0])
texts = [t for _, t, _ in rows]
digest = content_hash(texts)
if not force and state_matches(
extra, version=vocab.version, model=model, digest=digest
):
stats["skipped_state"] += 1
if progress:
progress(key, None, "unchanged")
continue
result = tag_comment(
key,
texts,
[v for _, _, v in rows],
vocab=vocab,
card_vecs=card_vecs,
judge=judge,
top=top,
max_tags=max_tags,
min_confidence=min_confidence,
)
stats["tagged"] += 1
if not dry_run:
store.merge_extra(
key,
{STATE_KEY: state_payload(result, version=vocab.version, model=model)},
)
added, _removed = apply_tags(store, key, result.slugs, vocab=vocab)
stats["tags_written"] += added
if progress:
progress(key, result, "tagged")
return stats
# ── real collaborators (the CLI wires these) ──────────────────────────
def card_vectors(vocab: Vocab, embed: Embed) -> dict[str, list[float]]:
slugs = list(vocab.cards)
vecs = embed([vocab.cards[s] for s in slugs])
return dict(zip(slugs, vecs))
def make_judge(
cfg: Any, pool: Any, *, post: Callable[..., dict] | None = None
) -> Judge:
"""A yes/no judge on the largest live host, ``temperature 0`` — the
same route ``llm.classify.closed_vocab_classifier`` takes."""
from llm.classify import _default_post
from llm.pool import pick_model
send = post or _default_post
pool.check(cfg.instruct_model)
def judge(theme: Theme, evidence: str) -> bool | None:
with pool.acquire_generation() as host:
model = pick_model(cfg, pool, host)
data = send(
f"{host}/api/chat",
json={
"model": model,
"messages": judge_prompt(theme, evidence),
"stream": False,
"options": {"num_ctx": cfg.chat_num_ctx, "temperature": 0},
},
)
reply = data.get("message", {}).get("content", "")
verdict = parse_yes_no(reply)
if verdict is None:
log.info("judge: unparseable reply %r for %s", reply[:60], theme.slug)
return verdict
return judge
def pg_chunks(engine: Any, collection: str = "comments") -> ChunkLoader:
"""``load_chunks(key)`` over ``langchain_pg_embedding`` — the text and
vector of every chunk indexed for the item, in ``seq`` order."""
from sqlalchemy import text as _text
sql = _text(
"SELECT (e.cmetadata->>'seq')::int AS seq, e.document, e.embedding::text "
"FROM langchain_pg_embedding e JOIN langchain_pg_collection c ON c.uuid = e.collection_id "
"WHERE c.name = :collection AND e.cmetadata->>'item_key' = :key ORDER BY seq"
)
def load(key: str) -> list[tuple[int, str, list[float]]]:
with engine.begin() as con:
rows = con.execute(sql, {"collection": collection, "key": key}).fetchall()
return [(int(seq or 0), doc or "", json.loads(vec)) for seq, doc, vec in rows]
return load

View File

@@ -0,0 +1,103 @@
"""stack llm tag / vocab (#574, #575)."""
from __future__ import annotations
import pytest
from typer.testing import CliRunner
import cli.llm as llm_cli
from bib.item import Source
from bib.store import Store
from cli import app
from llm.vocab import parse
runner = CliRunner()
VOCAB = parse(
"version: 1\nthemes:\n - slug: telehealth\n definition: Telehealth.\n synonyms: [telehealth]\n"
" - slug: drugs\n definition: Drugs.\n synonyms: [ASP]\n"
)
@pytest.fixture
def rt(monkeypatch, tmp_path):
store = Store(":memory:", storage_dir=tmp_path / "st")
key = store.upsert(
Source(
title="c1",
url="https://www.regulations.gov/comment/CMS-2026-2377-1",
date_published="2026-09-01",
),
tags=["reg-docket:CMS-2026-2377", "source:regulations-gov"],
)
judged = []
def judge(theme, evidence):
judged.append(theme.slug)
return theme.slug == "telehealth"
runtime = {
"store": store,
"vocab": VOCAB,
"load_chunks": lambda k: (
[(0, "telehealth text", [1.0, 0.0])] if k == key else []
),
"card_vecs": {"telehealth": [1.0, 0.0], "drugs": [0.0, 1.0]},
"judge": judge,
"model": "m",
}
monkeypatch.setattr(llm_cli, "_tagging_runtime", lambda cfg: runtime)
monkeypatch.setattr("llm.config.load", lambda: object())
try:
yield store, key, judged
finally:
store.close()
class TestTag:
def test_tags_and_reports(self, rt):
store, key, judged = rt
res = runner.invoke(
app, ["llm", "tag", "--docket", "CMS-2026-2377", "--verbose"]
)
assert res.exit_code == 0, res.output
assert f"{key} telehealth" in res.output
assert "seen=1 tagged=1 unchanged=0 no_chunks=0 tags_written=1" in res.output
assert (
"llm:telehealth" in store.get(key).tags
and "llm:drugs" not in store.get(key).tags
)
assert judged == ["telehealth", "drugs"]
res = runner.invoke(app, ["llm", "tag", "--docket", "CMS-2026-2377"])
assert "unchanged=1" in res.output
def test_dry_run_prints_but_writes_nothing(self, rt):
store, key, _ = rt
res = runner.invoke(
app, ["llm", "tag", "--docket", "CMS-2026-2377", "--dry-run"]
)
assert res.exit_code == 0, res.output
assert "dry run" in res.output and f"{key} telehealth" in res.output
assert not any(t.startswith("llm:") for t in store.get(key).tags)
def test_docket_is_required(self):
assert runner.invoke(app, ["llm", "tag"]).exit_code != 0
class TestVocab:
def test_lists_and_shows(self):
res = runner.invoke(app, ["llm", "vocab"])
assert (
res.exit_code == 0
and "vocabulary v" in res.output
and "telehealth" in res.output
)
res = runner.invoke(app, ["llm", "vocab", "--slug", "telehealth"])
assert res.exit_code == 0 and "llm:telehealth" in res.output
assert runner.invoke(app, ["llm", "vocab", "--slug", "nope"]).exit_code != 0
def test_invalid_file(self, tmp_path):
p = tmp_path / "bad.yaml"
p.write_text("version: 1\nthemes: []\n")
res = runner.invoke(app, ["llm", "vocab", "--path", str(p)])
assert res.exit_code == 1 and "invalid vocabulary" in res.output

390
tests/llm/test_tagger.py Normal file
View File

@@ -0,0 +1,390 @@
"""llm.tagger — shortlist, judge, state, write-back and the docket runner (#575/#576)."""
from __future__ import annotations
import json
import pytest
from bib.item import Source
from bib.store import Store
from llm.tagger import (
Candidate,
TagResult,
apply_tags,
card_vectors,
content_hash,
cosine,
docket_keys,
judge_prompt,
make_judge,
parse_yes_no,
run,
shortlist,
state_matches,
state_payload,
tag_comment,
)
from llm.vocab import parse
VOCAB = parse(
"""
version: 2
themes:
- slug: telehealth
label: Telehealth
definition: Medicare telehealth services.
synonyms: [telehealth]
- slug: care-management
label: Care management
definition: Monthly care management codes.
synonyms: [CCM]
- slug: drugs
label: Part B drugs
definition: ASP drug payment.
synonyms: [ASP]
"""
)
CARDS = {
"telehealth": [1.0, 0.0, 0.0],
"care-management": [0.0, 1.0, 0.0],
"drugs": [0.0, 0.0, 1.0],
}
class TestPure:
def test_cosine(self):
assert cosine([1, 0], [1, 0]) == pytest.approx(1.0)
assert cosine([1, 0], [0, 1]) == pytest.approx(0.0)
assert cosine([0, 0], [1, 1]) == 0.0
def test_shortlist_best_chunk_per_theme(self):
chunks = [[0.9, 0.1, 0.0], [0.1, 0.9, 0.0]]
out = shortlist(chunks, CARDS, top=2)
assert [c.slug for c in out] == ["care-management", "telehealth"] or [
c.slug for c in out
] == ["telehealth", "care-management"]
by = {c.slug: c for c in out}
assert by["telehealth"].chunk_idx == 0 and by["care-management"].chunk_idx == 1
assert all(isinstance(c, Candidate) and 0 < c.score <= 1 for c in out)
assert shortlist([], CARDS) == []
def test_judge_prompt_and_parse(self):
msgs = judge_prompt(
VOCAB.get("telehealth"), "We support audio-only telehealth."
)
assert msgs[0]["role"] == "system" and "Telehealth" in msgs[1]["content"]
assert "audio-only telehealth" in msgs[1]["content"]
assert parse_yes_no("Yes.") is True and parse_yes_no(" no\n") is False
assert parse_yes_no("Maybe yes") is None and parse_yes_no("") is None
def test_content_hash_is_order_sensitive(self):
assert content_hash(["a", "b"]) != content_hash(["b", "a"])
assert len(content_hash(["a"])) == 16
def test_tag_comment_accepts_judged_yes_up_to_max(self):
chunks = ["telehealth text", "ccm text", "asp text"]
vecs = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
seen = []
def judge(theme, evidence):
seen.append((theme.slug, evidence))
return theme.slug != "drugs"
r = tag_comment(
"K",
chunks,
vecs,
vocab=VOCAB,
card_vecs=CARDS,
judge=judge,
top=3,
max_tags=1,
)
assert isinstance(r, TagResult) and len(r.tags) == 1
assert r.judged["drugs"] is False and any(v for v in r.judged.values())
assert (
"telehealth",
"telehealth text",
) in seen # evidence is the scoring chunk
assert r.content_hash == content_hash(chunks)
def test_min_confidence_filters_weak_matches(self):
chunks = ["x"]
vecs = [[0.6, 0.8, 0.0]] # cosine 0.6 with telehealth, 0.8 with care-management
r = tag_comment(
"K",
chunks,
vecs,
vocab=VOCAB,
card_vecs=CARDS,
judge=lambda t, e: True,
min_confidence=0.7,
)
assert r.slugs == ["care-management"]
def test_state(self):
r = tag_comment(
"K",
["t"],
[[1.0, 0, 0]],
vocab=VOCAB,
card_vecs=CARDS,
judge=lambda t, e: True,
max_tags=1,
)
payload = state_payload(r, version=2, model="m")
assert payload["tags"] == {
"telehealth": {"confidence": 1.0, "evidence_chunk": 0}
}
assert payload["shortlist"][0]["verdict"] is True and "tagged_at" in payload
assert state_matches(
{"llm_tags": payload}, version=2, model="m", digest=r.content_hash
)
assert not state_matches(
{"llm_tags": payload}, version=3, model="m", digest=r.content_hash
)
assert not state_matches({}, version=2, model="m", digest=r.content_hash)
def test_card_vectors_embeds_every_card_once(self):
calls = []
def embed(texts):
calls.append(list(texts))
return [[float(i)] for i in range(len(texts))]
cv = card_vectors(VOCAB, embed)
assert set(cv) == {"telehealth", "care-management", "drugs"} and len(calls) == 1
assert "Medicare telehealth services" in calls[0][0]
class TestJudge:
def test_make_judge_posts_to_largest_host_at_temperature_zero(self):
class Pool:
def check(self, model):
return ["http://h"]
def vram(self, host):
return 0.0
def serves(self, host, model):
return False
class _cm:
def __enter__(self):
return "http://h"
def __exit__(self, *a):
return False
def acquire_generation(self):
return self._cm()
class Cfg:
instruct_model = "m"
instruct_model_large = ""
large_min_vram_gb = 20.0
chat_num_ctx = 4096
sent = {}
def post(url, json=None):
sent["url"], sent["json"] = url, json
return {"message": {"content": "Yes"}}
judge = make_judge(Cfg(), Pool(), post=post)
assert judge(VOCAB.get("drugs"), "ASP add-on") is True
assert (
sent["url"] == "http://h/api/chat"
and sent["json"]["options"]["temperature"] == 0
)
assert sent["json"]["model"] == "m"
def _store(tmp_path):
s = Store(":memory:", storage_dir=tmp_path / "st")
keys = {}
for i, (cid, date) in enumerate(
[
("CMS-2026-2377-1", "2026-09-01"),
("CMS-2026-2377-2", "2026-09-05"),
("CMS-2025-0304-9", "2025-09-01"),
]
):
docket = cid.rsplit("-", 1)[0]
keys[cid] = s.upsert(
Source(
title=cid,
url=f"https://www.regulations.gov/comment/{cid}",
date_published=date,
),
tags=[f"reg-docket:{docket}", "source:regulations-gov", "topic:palliative"],
)
return s, keys
class TestBibSide:
def test_docket_keys_newest_first(self, tmp_path):
s, keys = _store(tmp_path)
assert docket_keys(s, "CMS-2026-2377") == [
keys["CMS-2026-2377-2"],
keys["CMS-2026-2377-1"],
]
assert docket_keys(s, "CMS-2026-2377", limit=1) == [keys["CMS-2026-2377-2"]]
assert docket_keys(s, "NOPE") == []
s.close()
def test_apply_tags_replaces_only_llm_tags(self, tmp_path):
s, keys = _store(tmp_path)
k = keys["CMS-2026-2377-1"]
s.add_tags(k, ["llm:drugs"])
added, removed = apply_tags(
s, k, ["telehealth", "care-management"], vocab=VOCAB
)
assert (added, removed) == (2, 1)
tags = set(s.get(k).tags)
assert {
"llm:telehealth",
"llm:care-management",
"topic:palliative",
"reg-docket:CMS-2026-2377",
} <= tags
assert "llm:drugs" not in tags
assert apply_tags(s, k, ["telehealth", "care-management"], vocab=VOCAB) == (
0,
0,
) # idempotent
assert (
apply_tags(s, k, [], vocab=VOCAB) == (0, 2)
and "topic:palliative" in s.get(k).tags
)
s.close()
def test_merge_extra(self, tmp_path):
s, keys = _store(tmp_path)
k = keys["CMS-2026-2377-1"]
s.merge_extra(k, {"a": 1})
s.merge_extra(k, {"b": {"x": 2}})
s.merge_extra("ZZZZZZZZ", {"c": 3}) # unknown key: no-op, no error
extra = json.loads(
s._con()
.execute("SELECT extra_json FROM items WHERE key=?", (k,))
.fetchone()[0]
) # noqa: SLF001
assert extra["a"] == 1 and extra["b"] == {"x": 2}
s.close()
class TestRun:
def _chunks(self, keys):
def load(key):
if key == keys["CMS-2026-2377-1"]:
return [
(1, "ccm text", [0.0, 1.0, 0.0]),
(0, "telehealth text", [1.0, 0.0, 0.0]),
]
if key == keys["CMS-2026-2377-2"]:
return [] # not indexed yet
return [(0, "asp", [0.0, 0.0, 1.0])]
return load
def test_tags_persist_state_and_resume(self, tmp_path):
s, keys = _store(tmp_path)
events = []
judge_calls = []
def judge(theme, evidence):
judge_calls.append(theme.slug)
return True
stats = run(
s,
vocab=VOCAB,
docket="CMS-2026-2377",
load_chunks=self._chunks(keys),
card_vecs=CARDS,
judge=judge,
model="m",
top=2,
max_tags=2,
progress=lambda k, r, why: events.append((k, why)),
)
assert stats == {
"seen": 2,
"tagged": 1,
"skipped_state": 0,
"no_chunks": 1,
"tags_written": 2,
}
k = keys["CMS-2026-2377-1"]
assert {"llm:telehealth", "llm:care-management"} <= set(s.get(k).tags)
extra = json.loads(
s._con()
.execute("SELECT extra_json FROM items WHERE key=?", (k,))
.fetchone()[0]
) # noqa: SLF001
assert extra["llm_tags"]["version"] == 2 and extra["llm_tags"]["model"] == "m"
assert set(extra["llm_tags"]["tags"]) == {"telehealth", "care-management"}
assert (keys["CMS-2026-2377-2"], "no-chunks") in events
# second run: unchanged → skipped without judging
judge_calls.clear()
stats2 = run(
s,
vocab=VOCAB,
docket="CMS-2026-2377",
load_chunks=self._chunks(keys),
card_vecs=CARDS,
judge=judge,
model="m",
top=2,
)
assert (
stats2["skipped_state"] == 1 and stats2["tagged"] == 0 and judge_calls == []
)
# a new model re-tags; --force re-tags
assert (
run(
s,
vocab=VOCAB,
docket="CMS-2026-2377",
load_chunks=self._chunks(keys),
card_vecs=CARDS,
judge=judge,
model="m2",
top=2,
)["tagged"]
== 1
)
assert (
run(
s,
vocab=VOCAB,
docket="CMS-2026-2377",
load_chunks=self._chunks(keys),
card_vecs=CARDS,
judge=judge,
model="m2",
top=2,
force=True,
)["tagged"]
== 1
)
s.close()
def test_dry_run_writes_nothing(self, tmp_path):
s, keys = _store(tmp_path)
stats = run(
s,
vocab=VOCAB,
docket="CMS-2026-2377",
load_chunks=self._chunks(keys),
card_vecs=CARDS,
judge=lambda t, e: True,
model="m",
dry_run=True,
)
assert stats["tagged"] == 1 and stats["tags_written"] == 0
k = keys["CMS-2026-2377-1"]
assert not any(t.startswith("llm:") for t in s.get(k).tags)
s.close()