homelab/stack
1
0
Fork 0
Code Issues 275 Pull Requests 2 Actions Packages Projects 1 Releases Wiki Activity

feat(llm): rule-text source + rules collection CLI (refs #595)

Browse Source
This commit is contained in:
kert
2026-08-13 12:19:26 -04:00
parent ab090e2840
commit ad6d3f05c1
4 changed files with 176 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/cli/llm.py
View File

@@ -8,20 +8,21 @@ app = typer.Typer(no_args_is_help=True)
@app.command()
def index(
collection: str = typer.Option(
"comments", help="Which collection: comments | corpus."
"comments", help="Which collection: comments | corpus | rules."
),
docket: str = typer.Option("", help="Limit comments to one docket id."),
key: list[str] = typer.Option([], "--key", help="Limit rules to these item keys."),
force: bool = typer.Option(False, help="Re-embed even when unchanged."),
limit: int = typer.Option(0, help="Stop after N docs (0 = all)."),
) -> None:
"""Embed comments/corpus into pgvector (incremental, resumable)."""
"""Embed comments/corpus/rules into pgvector (incremental, resumable)."""
import itertools
from conf.connect import bib
from llm import config as llm_config
from llm.index import index_docs
from llm.pool import HostPool
from llm.source import iter_comment_docs, iter_corpus_docs
from llm.source import iter_comment_docs, iter_corpus_docs, iter_rule_docs
cfg = llm_config.load()
store = bib()
@@ -29,8 +30,10 @@ def index(
docs = iter_comment_docs(store, docket=docket)
elif collection == "corpus":
docs = iter_corpus_docs(store)
elif collection == "rules":
docs = iter_rule_docs(store, keys=tuple(key))
else:
raise typer.BadParameter("collection must be 'comments' or 'corpus'")
raise typer.BadParameter("collection must be 'comments', 'corpus' or 'rules'")
if limit:
docs = itertools.islice(docs, limit)
stats = index_docs(

63
src/llm/source.py
View File

@@ -123,6 +123,69 @@ def _attachment_text(store: Store, item_key: str) -> str:
return "\n\n".join(parts)
def _strip_html_wrapper(text: str) -> str:
"""Drop the FR ``.txt`` attachment's HTML wrapper, if present.
FR TXT attachments are HTML-wrapped: ``<html><head>...<body><pre>``
then the plain text, closed with ``</pre></body></html>``. Strip
everything through the first ``<pre>`` and any trailing close tags;
text without that wrapper passes through unchanged.
"""
idx = text.find("<pre>")
if idx == -1:
return text
body = text[idx + len("<pre>") :]
stripped = body.rstrip()
for suffix in ("</pre></body></html>", "</pre>"):
if stripped.endswith(suffix):
return stripped[: -len(suffix)]
return body
def _rule_text(store: Store, item_key: str) -> str:
"""TXT attachment (HTML wrapper stripped) preferred; else PDF-extract."""
rows = (
store._con()
.execute(
"SELECT a.storage_path FROM attachments a "
"JOIN items i ON i.id = a.item_id WHERE i.key = ?",
(item_key,),
)
.fetchall()
)
for (storage_path,) in rows:
path = Path(storage_path)
if path.suffix.lower() == ".txt" and path.exists():
return _strip_html_wrapper(path.read_text())
return _attachment_text(store, item_key)
def iter_rule_docs(
store: Store, *, keys: tuple[str, ...] = (), tag: str = ""
) -> Iterator[Doc]:
"""One Doc per FR rule item: TXT attachment preferred, PDF-extract fallback."""
for item in store.list_items(item_type="rule", tag=tag):
if keys and item.key not in keys:
continue
text = _rule_text(store, item.key)
if not text.strip():
continue
cms_rule = next(
(t.split(":", 1)[1] for t in item.tags if t.startswith("cms-rule:")), ""
)
yield Doc(
key=item.key,
text=text,
metadata={
"doctype": "rule",
"cms_rule_id": cms_rule,
"fr_document_number": item.document_number or "",
"year": _year_of(store, item.key),
"item_key": item.key,
},
)
def iter_corpus_docs(store: Store, *, tag: str = "") -> Iterator[Doc]:
"""Every non-comment item; attachment text + abstract."""
for item in store.list_items(tag=tag):

31
tests/cli/test_llm_exercise.py
View File

@@ -71,6 +71,35 @@ class TestIndexCorpus:
assert "indexed=3 skipped=1 chunks=7" in result.output
class TestIndexRules:
@patch("llm.source.iter_rule_docs")
@patch("llm.pool.HostPool.from_config")
@patch("llm.index.index_docs")
@patch("conf.connect.bib")
@patch("llm.config.load")
def test_rules_collection(
self, mock_load, mock_bib, mock_index_docs, mock_from_config, mock_iter
):
cfg = MagicMock()
mock_load.return_value = cfg
store = MagicMock()
mock_bib.return_value = store
mock_iter.return_value = iter(["doc1"])
mock_from_config.return_value = MagicMock()
mock_index_docs.return_value = _STATS
result = runner.invoke(
app,
["index", "--collection", "rules", "--key", "abc123", "--key", "def456"],
)
assert result.exit_code == 0
mock_iter.assert_called_once_with(store, keys=("abc123", "def456"))
kwargs = mock_index_docs.call_args.kwargs
assert kwargs["collection"] == "rules"
assert "indexed=3 skipped=1 chunks=7" in result.output
class TestIndexBadCollection:
@patch("conf.connect.bib")
@patch("llm.config.load")
@@ -81,7 +110,7 @@ class TestIndexBadCollection:
result = runner.invoke(app, ["index", "--collection", "bogus"])
assert result.exit_code == 2
assert "collection must be 'comments' or 'corpus'" in result.output
assert "collection must be 'comments', 'corpus' or 'rules'" in result.output
class TestIndexLimit:

77
tests/llm/test_source.py
View File

@@ -1,14 +1,17 @@
"""llm.source — comment + corpus Doc iterators."""
from unittest.mock import patch
import pytest
from bib.item import Item
from bib.item import Item, Rule
from bib.store import Store
from llm.source import (
_default_root,
comment_key_map,
iter_comment_docs,
iter_corpus_docs,
iter_rule_docs,
)
DOCKET = "CMS-2019-0111"
@@ -104,6 +107,78 @@ class TestCorpusDocs:
assert key not in {d.key for d in iter_corpus_docs(store)}
class TestRuleDocs:
@pytest.fixture
def rule_key(self, store):
key = store.create(
Rule(
title="CY2026 PFS Proposed Rule",
document_number="2025-13271",
)
)
for tag in ("cms-rule:CMS-1832-P", "year:2026"):
store.add_tag(key, tag)
return key
def test_txt_attachment_yields_doc_with_metadata(self, store, rule_key, tmp_path):
txt = tmp_path / "2025-13271.txt"
txt.write_text(
"<html><head><title>x</title></head><body><pre>\n"
"The Secretary proposes to amend 42 CFR part 414.\n"
"</pre></body></html>"
)
store.attach_file(rule_key, txt)
docs = list(iter_rule_docs(store))
assert len(docs) == 1
doc = docs[0]
assert doc.key == rule_key
assert "The Secretary proposes to amend 42 CFR part 414." in doc.text
assert "<html>" not in doc.text
assert "<pre>" not in doc.text
assert doc.metadata == {
"doctype": "rule",
"cms_rule_id": "CMS-1832-P",
"fr_document_number": "2025-13271",
"year": "2026",
"item_key": rule_key,
}
def test_pdf_only_falls_back_to_extract_attachment(self, store, rule_key, tmp_path):
pdf = tmp_path / "rule.pdf"
pdf.write_bytes(b"%PDF-1.4 fake")
store.attach_file(rule_key, pdf)
with patch("rex.comments.combine.extract_attachment") as mock_extract:
from rex.comments.extract import ExtractResult
mock_extract.return_value = ExtractResult(
text="Extracted PDF body.", status="ok", chars=20
)
docs = list(iter_rule_docs(store))
assert len(docs) == 1
assert docs[0].text == "Extracted PDF body."
def test_keys_filters_to_named_items(self, store, rule_key, tmp_path):
other_key = store.create(Rule(title="Other rule", document_number="2025-00001"))
store.add_tag(other_key, "cms-rule:CMS-9999-P")
for key in (rule_key, other_key):
txt = tmp_path / f"{key}.txt"
txt.write_text("<pre>Some rule text.</pre>")
store.attach_file(key, txt)
docs = list(iter_rule_docs(store, keys=(rule_key,)))
assert [d.key for d in docs] == [rule_key]
def test_non_rule_items_not_yielded(self, store):
assert all(d.metadata["doctype"] == "rule" for d in iter_rule_docs(store))
# the "report" comment item from the `store` fixture is never yielded
assert store._comment_key not in {d.key for d in iter_rule_docs(store)}
class TestDefaultRoot:
def test_default_root_under_state_comments(self):
from conf import ROOT