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

src/cli/llm.py

@@ -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(

src/llm/source.py

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

tests/cli/test_llm_exercise.py

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

tests/llm/test_source.py

@@ -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