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:
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack api serve
|
title: stack api serve
|
||||||
sidebar_position: 62
|
sidebar_position: 63
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack api serve`
|
# `stack api serve`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack api
|
title: stack api
|
||||||
sidebar_position: 61
|
sidebar_position: 62
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack api`
|
# `stack api`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack db comment
|
title: stack db comment
|
||||||
sidebar_position: 55
|
sidebar_position: 56
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack db comment`
|
# `stack db comment`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack db inspect
|
title: stack db inspect
|
||||||
sidebar_position: 56
|
sidebar_position: 57
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack db inspect`
|
# `stack db inspect`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack db
|
title: stack db
|
||||||
sidebar_position: 54
|
sidebar_position: 55
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack db`
|
# `stack db`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack docs build
|
title: stack docs build
|
||||||
sidebar_position: 58
|
sidebar_position: 59
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack docs build`
|
# `stack docs build`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack docs generate
|
title: stack docs generate
|
||||||
sidebar_position: 60
|
sidebar_position: 61
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack docs generate`
|
# `stack docs generate`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack docs serve
|
title: stack docs serve
|
||||||
sidebar_position: 59
|
sidebar_position: 60
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack docs serve`
|
# `stack docs serve`
|
||||||
|
|||||||
@@ -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
38
docs/docs/cli/llm-tag.md
Normal 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. │
|
||||||
|
╰──────────────────────────────────────────────────────────────────────────────╯
|
||||||
|
```
|
||||||
@@ -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. │
|
||||||
╰──────────────────────────────────────────────────────────────────────────────╯
|
╰──────────────────────────────────────────────────────────────────────────────╯
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack mail dns
|
title: stack mail dns
|
||||||
sidebar_position: 101
|
sidebar_position: 102
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack mail dns`
|
# `stack mail dns`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack mail down
|
title: stack mail down
|
||||||
sidebar_position: 99
|
sidebar_position: 100
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack mail down`
|
# `stack mail down`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack mail provision
|
title: stack mail provision
|
||||||
sidebar_position: 97
|
sidebar_position: 98
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack mail provision`
|
# `stack mail provision`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack mail status
|
title: stack mail status
|
||||||
sidebar_position: 100
|
sidebar_position: 101
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack mail status`
|
# `stack mail status`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack mail up
|
title: stack mail up
|
||||||
sidebar_position: 98
|
sidebar_position: 99
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack mail up`
|
# `stack mail up`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack mail
|
title: stack mail
|
||||||
sidebar_position: 96
|
sidebar_position: 97
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack mail`
|
# `stack mail`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack perf show
|
title: stack perf show
|
||||||
sidebar_position: 64
|
sidebar_position: 65
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack perf show`
|
# `stack perf show`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack perf
|
title: stack perf
|
||||||
sidebar_position: 63
|
sidebar_position: 64
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack perf`
|
# `stack perf`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs elements
|
title: stack pfs elements
|
||||||
sidebar_position: 70
|
sidebar_position: 71
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs elements`
|
# `stack pfs elements`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs exposure
|
title: stack pfs exposure
|
||||||
sidebar_position: 75
|
sidebar_position: 76
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs exposure`
|
# `stack pfs exposure`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs families
|
title: stack pfs families
|
||||||
sidebar_position: 72
|
sidebar_position: 73
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs families`
|
# `stack pfs families`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs guidance
|
title: stack pfs guidance
|
||||||
sidebar_position: 73
|
sidebar_position: 74
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs guidance`
|
# `stack pfs guidance`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs lineage
|
title: stack pfs lineage
|
||||||
sidebar_position: 71
|
sidebar_position: 72
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs lineage`
|
# `stack pfs lineage`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs reaction
|
title: stack pfs reaction
|
||||||
sidebar_position: 74
|
sidebar_position: 75
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs reaction`
|
# `stack pfs reaction`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs review
|
title: stack pfs review
|
||||||
sidebar_position: 77
|
sidebar_position: 78
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs review`
|
# `stack pfs review`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs utilization
|
title: stack pfs utilization
|
||||||
sidebar_position: 76
|
sidebar_position: 77
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs utilization`
|
# `stack pfs utilization`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack pfs
|
title: stack pfs
|
||||||
sidebar_position: 69
|
sidebar_position: 70
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack pfs`
|
# `stack pfs`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma eligible
|
title: stack prisma eligible
|
||||||
sidebar_position: 90
|
sidebar_position: 91
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma eligible`
|
# `stack prisma eligible`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma export
|
title: stack prisma export
|
||||||
sidebar_position: 87
|
sidebar_position: 88
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma export`
|
# `stack prisma export`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma extract
|
title: stack prisma extract
|
||||||
sidebar_position: 91
|
sidebar_position: 92
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma extract`
|
# `stack prisma extract`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma fetch
|
title: stack prisma fetch
|
||||||
sidebar_position: 93
|
sidebar_position: 94
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma fetch`
|
# `stack prisma fetch`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma flow
|
title: stack prisma flow
|
||||||
sidebar_position: 92
|
sidebar_position: 93
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma flow`
|
# `stack prisma flow`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma init
|
title: stack prisma init
|
||||||
sidebar_position: 86
|
sidebar_position: 87
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma init`
|
# `stack prisma init`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma run
|
title: stack prisma run
|
||||||
sidebar_position: 94
|
sidebar_position: 95
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma run`
|
# `stack prisma run`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma screen
|
title: stack prisma screen
|
||||||
sidebar_position: 89
|
sidebar_position: 90
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma screen`
|
# `stack prisma screen`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma vpn
|
title: stack prisma vpn
|
||||||
sidebar_position: 95
|
sidebar_position: 96
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma vpn`
|
# `stack prisma vpn`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack prisma
|
title: stack prisma
|
||||||
sidebar_position: 85
|
sidebar_position: 86
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack prisma`
|
# `stack prisma`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack rec list
|
title: stack rec list
|
||||||
sidebar_position: 66
|
sidebar_position: 67
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack rec list`
|
# `stack rec list`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack rec opps
|
title: stack rec opps
|
||||||
sidebar_position: 68
|
sidebar_position: 69
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack rec opps`
|
# `stack rec opps`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack rec pfs
|
title: stack rec pfs
|
||||||
sidebar_position: 67
|
sidebar_position: 68
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack rec pfs`
|
# `stack rec pfs`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack rec
|
title: stack rec
|
||||||
sidebar_position: 65
|
sidebar_position: 66
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack rec`
|
# `stack rec`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -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`
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
---
|
---
|
||||||
title: stack zot
|
title: stack zot
|
||||||
sidebar_position: 79
|
sidebar_position: 80
|
||||||
---
|
---
|
||||||
|
|
||||||
# `stack zot`
|
# `stack zot`
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
102
src/cli/llm.py
102
src/cli/llm.py
@@ -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
369
src/llm/tagger.py
Normal 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
|
||||||
103
tests/cli/test_llm_tag_cli.py
Normal file
103
tests/cli/test_llm_tag_cli.py
Normal 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
390
tests/llm/test_tagger.py
Normal 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()
|
||||||
Reference in New Issue
Block a user