feat(llm): comment + corpus Doc sources with abstract fallback (refs #566)

src/llm/source.py

@@ -0,0 +1,143 @@
+"""Doc iterators over the bib store + comment extraction tree.
+
+Comments: text prefers the #253 extraction output
+(``.state/comments/<docket>/<comment_id>/combined.md`` body via
+``rex.comments.combine.parse_combined``); falls back to the bib
+``abstract`` for attachment-less comments. Corpus: every non-comment item;
+attachment text extracted with the same rex extractors, joined with the
+abstract.
+
+Docket is derived from the comment's URL (``.../comment/<docket>-<seq>``)
+rather than a ``docket:`` tag — the live store tags comments with
+``reg-docket:<docket>`` (and a separate, unrelated empty ``docket:`` tag),
+so parsing the comment id is the reliable path to a comment's docket.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterator
+
+from bib.store import Store
+from llm.chunk import Doc
+
+_COMMENT_URL_PREFIX = "https://www.regulations.gov/comment/"
+
+
+def _default_root() -> Path:
+    from conf import ROOT
+
+    return ROOT / ".state" / "comments"
+
+
+def _year_of(store: Store, item_key: str) -> str:
+    row = (
+        store._con()
+        .execute(
+            "SELECT t.name FROM tags t "
+            "JOIN item_tags it ON it.tag_id = t.id "
+            "JOIN items i ON i.id = it.item_id "
+            "WHERE i.key = ? AND t.name LIKE 'year:%'",
+            (item_key,),
+        )
+        .fetchone()
+    )
+    return row[0].split(":", 1)[1] if row else ""
+
+
+def _comment_rows(store: Store, docket: str = "") -> list[tuple[str, str, str]]:
+    """(docket, comment_id, item key) for every comment matching *docket*.
+
+    *docket* filters via the URL (``.../comment/<docket>-%``); an empty
+    docket matches every comment and the docket is recovered from each
+    comment_id's ``<docket>-<seq>`` suffix.
+    """
+    pattern = (
+        f"{_COMMENT_URL_PREFIX}{docket}-%" if docket else f"{_COMMENT_URL_PREFIX}%"
+    )
+    rows = (
+        store._con()
+        .execute("SELECT i.key, i.url FROM items i WHERE i.url LIKE ?", (pattern,))
+        .fetchall()
+    )
+    out = []
+    for key, url in rows:
+        comment_id = url.rsplit("/", 1)[-1]
+        dk = docket or comment_id.rsplit("-", 1)[0]
+        out.append((dk, comment_id, key))
+    return out
+
+
+def comment_key_map(store: Store, docket: str) -> dict[str, tuple[str, str]]:
+    """comment_id -> (bib item key, year) for one docket."""
+    return {
+        comment_id: (key, _year_of(store, key))
+        for _, comment_id, key in _comment_rows(store, docket)
+    }
+
+
+def iter_comment_docs(
+    store: Store, *, docket: str = "", root: Path | None = None
+) -> Iterator[Doc]:
+    """One Doc per comment: extraction body, else abstract fallback."""
+    from rex.comments.combine import parse_combined
+
+    root = root if root is not None else _default_root()
+    for dk, comment_id, key in sorted(_comment_rows(store, docket)):
+        meta = {
+            "docket": dk,
+            "comment_id": comment_id,
+            "doctype": "comment",
+            "year": _year_of(store, key),
+        }
+        combined = root / dk / comment_id / "combined.md"
+        if combined.exists():
+            _, body = parse_combined(combined.read_text())
+            if body.strip():
+                yield Doc(key=key, text=body, metadata=meta)
+                continue
+        item = store.get(key)
+        if item.abstract.strip():
+            yield Doc(key=key, text=item.abstract, metadata=meta)
+
+
+def _attachment_text(store: Store, item_key: str) -> str:
+    from rex.comments.combine import extract_attachment
+
+    rows = (
+        store._con()
+        .execute(
+            "SELECT a.storage_path FROM attachments a "
+            "JOIN items i ON i.id = a.item_id WHERE i.key = ?",
+            (item_key,),
+        )
+        .fetchall()
+    )
+    parts = []
+    for (storage_path,) in rows:
+        path = Path(storage_path)
+        if path.exists():
+            result = extract_attachment(path)
+            if result.status == "ok":
+                parts.append(result.text)
+    return "\n\n".join(parts)
+
+
+def iter_corpus_docs(store: Store, *, tag: str = "") -> Iterator[Doc]:
+    """Every non-comment item; attachment text + abstract."""
+    for item in store.list_items(tag=tag):
+        if "doctype:comment" in item.tags:
+            continue
+        text = "\n\n".join(
+            p for p in (_attachment_text(store, item.key), item.abstract) if p.strip()
+        )
+        if not text.strip():
+            continue
+        yield Doc(
+            key=item.key,
+            text=text,
+            metadata={
+                "doctype": item.item_type,
+                "year": _year_of(store, item.key),
+            },
+        )

tests/llm/test_source.py

@@ -0,0 +1,85 @@
+"""llm.source — comment + corpus Doc iterators."""
+
+import pytest
+
+from bib.item import Item
+from bib.store import Store
+from llm.source import comment_key_map, iter_comment_docs, iter_corpus_docs
+
+DOCKET = "CMS-2019-0111"
+CID = f"{DOCKET}-0042"
+
+COMBINED = f"""---
+comment_id: {CID}
+docket_id: {DOCKET}
+---
+
+We object to the E/M consolidation.
+"""
+
+
+@pytest.fixture
+def store():
+    s = Store(":memory:")
+    key = s.create(
+        Item(
+            item_type="report",
+            title="A comment",
+            url=f"https://www.regulations.gov/comment/{CID}",
+            abstract="Inline abstract body.",
+        )
+    )
+    for tag in (f"docket:{DOCKET}", "doctype:comment", "year:2019"):
+        s.add_tag(key, tag)
+    s._comment_key = key  # test convenience
+    return s
+
+
+@pytest.fixture
+def root(tmp_path):
+    d = tmp_path / DOCKET / CID
+    d.mkdir(parents=True)
+    (d / "combined.md").write_text(COMBINED)
+    return tmp_path
+
+
+class TestCommentDocs:
+    def test_extracted_comment_uses_combined_body(self, store, root):
+        docs = list(iter_comment_docs(store, docket=DOCKET, root=root))
+        assert len(docs) == 1
+        assert docs[0].key == store._comment_key
+        assert "E/M consolidation" in docs[0].text
+        assert docs[0].metadata == {
+            "docket": DOCKET,
+            "comment_id": CID,
+            "doctype": "comment",
+            "year": "2019",
+        }
+
+    def test_unextracted_comment_falls_back_to_abstract(self, store, tmp_path):
+        docs = list(iter_comment_docs(store, docket=DOCKET, root=tmp_path))
+        assert docs[0].text == "Inline abstract body."
+
+    def test_docket_filter_excludes_others(self, store, root):
+        assert list(iter_comment_docs(store, docket="CMS-2021-0119", root=root)) == []
+
+
+class TestCommentKeyMap:
+    def test_maps_comment_id_to_key_and_year(self, store):
+        mapping = comment_key_map(store, DOCKET)
+        assert mapping[CID] == (store._comment_key, "2019")
+
+
+class TestCorpusDocs:
+    def test_non_comment_item_with_abstract(self, store):
+        key = store.create(
+            Item(item_type="rule", title="Final rule", abstract="Rule text.")
+        )
+        store.add_tag(key, "year:2020")
+        docs = list(iter_corpus_docs(store))
+        assert [d.key for d in docs] == [key]
+        assert docs[0].text == "Rule text."
+        assert docs[0].metadata["doctype"] == "rule"
+
+    def test_comments_excluded_from_corpus(self, store):
+        assert all(d.metadata["doctype"] != "comment" for d in iter_corpus_docs(store))