stack/tests/llm/test_source.py

"""llm.source — comment + corpus Doc iterators."""

from unittest.mock import patch

import pytest

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"
CID = f"{DOCKET}-0042"

COMBINED = f"""---
comment_id: {CID}
docket_id: {DOCKET}
---

We object to the E/M consolidation.
"""


@pytest.fixture
def store(tmp_path):
    s = Store(":memory:", storage_dir=tmp_path / "storage")
    key = s.create(
        Item(
            item_type="report",
            title="A comment",
            url=f"https://www.regulations.gov/comment/{CID}",
            abstract="Inline abstract body.",
            date_published="2019-09-27",
        )
    )
    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",
            "kind": "comment",
            "year": "2019",
            "date": "2019-09-27",
            "title": "A comment",
        }
        assert docs[0].files == ()

    def test_files_lists_attachments_in_comment_dir(self, store, root):
        (root / DOCKET / CID / "attachment_1.pdf").write_bytes(b"%PDF")
        (root / DOCKET / CID / "attachment_1.pdf.md").write_text("sibling")
        docs = list(iter_comment_docs(store, docket=DOCKET, root=root))
        assert docs[0].files == (
            ("attachment_1.pdf", str(root / DOCKET / CID / "attachment_1.pdf")),
        )

    def test_newest_first(self, store, root):
        older = store.create(
            Item(
                item_type="report",
                title="Older",
                url=f"https://www.regulations.gov/comment/{DOCKET}-0001",
                abstract="old body",
                date_published="2018-01-01",
            )
        )
        keys = [d.key for d in iter_comment_docs(store, docket=DOCKET, root=root)]
        assert keys == [store._comment_key, older]

    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="report",
                title="Final rule",
                abstract="Rule text.",
                url="https://x.test/r",
                date_published="2020-11-02",
            )
        )
        for t in ("year:2020", "project:pfs"):
            store.add_tag(key, t)
        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": "report",
            "kind": "corpus",
            "year": "2020",
            "date": "2020-11-02",
            "title": "Final rule",
            "url": "https://x.test/r",
            "project": "pfs",
        }

    def test_comments_excluded_from_corpus(self, store):
        assert all(d.metadata["doctype"] != "comment" for d in iter_corpus_docs(store))

    def test_llm_skip_tag_excludes_item_from_corpus(self, store):
        key = store.create(
            Item(item_type="source", title="Scratch export", abstract="Draft notes.")
        )
        store.add_tag(key, "llm:skip")
        assert key not in {d.key for d in iter_corpus_docs(store)}

    def test_rule_items_never_enter_the_corpus(self, store):
        """FR rules live in the rules collection (anchored paragraphs,
        federalregister.gov links); their PDFs must not be chunked a
        second time as 'corpus'."""
        typed = store.create(
            Item(item_type="rule", title="A rule", abstract="Rule text.")
        )
        anchored = store.create(
            Item(item_type="source", title="Grabbed", abstract="Body.")
        )
        store._con().execute(  # noqa: SLF001
            "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) VALUES (?, 1, 1, 1, 'x')",
            (anchored,),
        )
        keys = {d.key for d in iter_corpus_docs(store)}
        assert typed not in keys and anchored not in keys

    def test_attachment_text_sectioned_and_listed_in_files(self, store, tmp_path):
        key = store.create(Item(item_type="source", title="Report with attachment"))
        store.add_tag(key, "year:2021")
        att = tmp_path / "letter.txt"
        att.write_text("Attachment body text " * 10)  # > 50 chars => status "ok"
        store.attach_file(key, att)
        docs = {d.key: d for d in iter_corpus_docs(store)}
        assert docs[key].text.startswith("## letter.txt\n\nAttachment body text")
        assert [name for name, _ in docs[key].files] == ["letter.txt"]

    def test_epub_attachment_text_included(self, store, tmp_path):
        """F2: an EPUB attachment (the AMA's CPT codebooks) is real book
        text for the corpus indexer, same as a PDF/DOCX/txt attachment."""
        import zipfile

        key = store.create(Item(item_type="source", title="CPT Professional 2024"))
        store.add_tag(key, "year:2024")
        epub = tmp_path / "CPT Professional 2024.epub"
        with zipfile.ZipFile(epub, "w") as zf:
            zf.writestr("mimetype", "application/epub+zip")
            zf.writestr(
                "OPS/Chapter01.xhtml",
                "<html><body><p>"
                + "Chronic care management services, first 20 minutes. " * 4
                + "</p></body></html>",
            )
        store.attach_file(key, epub)
        docs = {d.key: d for d in iter_corpus_docs(store)}
        assert "Chronic care management services" in docs[key].text
        assert [name for name, _ in docs[key].files] == [epub.name]

    def test_zotero_pdf_fallback_used_when_bib_has_no_attachments(
        self, store, tmp_path
    ):
        from llm.source import ZoteroPdfIndex

        key = store.create(
            Item(item_type="source", title="PubMed record", abstract="Abs.")
        )
        store.add_tag(key, "year:2024")
        pdf = tmp_path / "paper.pdf"
        pdf.write_bytes(b"%PDF-1.4 fake")
        zot = ZoteroPdfIndex({key: [pdf]})
        with patch("rex.comments.combine.extract_attachment") as mock_extract:
            from rex.comments.extract import ExtractResult

            mock_extract.return_value = ExtractResult(
                text="Paper body.", status="ok", chars=11
            )
            docs = {d.key: d for d in iter_corpus_docs(store, zotero=zot)}
        assert docs[key].text == "## paper.pdf\n\nPaper body.\n\nAbs."
        assert docs[key].files == (("paper.pdf", str(pdf)),)

    def test_empty_text_item_skipped(self, store):
        key = store.create(Item(item_type="rule", title="No content", abstract="   "))
        store.add_tag(key, "year:2022")
        assert key not in {d.key for d in iter_corpus_docs(store)}


class TestZoteroPdfIndex:
    def test_snapshot_maps_parent_key_to_storage_pdfs(self, tmp_path):
        import sqlite3

        from llm.source import ZoteroPdfIndex

        db = tmp_path / "zotero.sqlite"
        con = sqlite3.connect(db)
        con.executescript(
            """
            CREATE TABLE items (itemID INTEGER PRIMARY KEY, key TEXT);
            CREATE TABLE deletedItems (itemID INTEGER);
            CREATE TABLE itemAttachments (itemID INTEGER, parentItemID INTEGER, path TEXT);
            INSERT INTO items VALUES (1,'PARENTK1'),(2,'ATTKEY01'),(3,'ATTKEY02'),(4,'GONEKEY1');
            INSERT INTO itemAttachments VALUES (2,1,'storage:paper.pdf'),(3,1,'storage:notes.txt'),(4,1,'storage:gone.pdf');
            INSERT INTO deletedItems VALUES (4);
            """
        )
        con.commit()
        con.close()
        storage = tmp_path / "storage"
        (storage / "ATTKEY01").mkdir(parents=True)
        (storage / "ATTKEY01" / "paper.pdf").write_bytes(b"%PDF")
        idx = ZoteroPdfIndex.snapshot(db, storage, tmp_path / "snap")
        assert idx.pdfs_for("PARENTK1") == [storage / "ATTKEY01" / "paper.pdf"]
        assert idx.pdfs_for("NOPE") == []

    def test_snapshot_missing_db_is_empty_index(self, tmp_path):
        from llm.source import ZoteroPdfIndex

        idx = ZoteroPdfIndex.snapshot(
            tmp_path / "none.sqlite", tmp_path, tmp_path / "s"
        )
        assert idx.pdfs_for("X") == []


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.pop("rule_kind") in (
            "proposed",
            "final",
            "correction",
            "rule",
        )
        assert doc.metadata == {
            "doctype": "rule",
            "kind": "rule",
            "cms_rule_id": "CMS-1832-P",
            "fr_document_number": "2025-13271",
            "year": "2026",
            "date": "",
            "title": "CY2026 PFS Proposed Rule",
            "item_key": rule_key,
            "html_url": "",
            "fr_volume": "",
        }
        assert doc.paragraphs == ()

    def test_grabbed_rule_yields_paragraphs_from_fr_anchors(self, store, rule_key):
        from bib import frlink

        html = """
        <p id="p-1" data-page="100">First para.</p>
        <p id="p-2" data-page="101">Second para.</p>
        """
        meta = {
            "html_url": "https://fr.test/doc",
            "body_html_url": "https://fr.test/body",
            "start_page": 100,
            "end_page": 101,
            "volume": 91,
        }
        frlink.grab(store, rule_key, fetch=lambda _d: (meta, html))

        (doc,) = list(iter_rule_docs(store))

        assert doc.text == "First para.\n\nSecond para."
        assert [(p.p_id, p.page, p.ordinal) for p in doc.paragraphs] == [
            (1, 100, 1),
            (2, 101, 1),
        ]
        assert doc.metadata["html_url"] == "https://fr.test/doc"
        assert doc.metadata["fr_volume"] == "91"
        assert doc.metadata["kind"] == "rule"

    def test_inline_html_and_entities_stripped(self, store, rule_key, tmp_path):
        txt = tmp_path / "2025-13271.txt"
        txt.write_text(
            "<html><head><title>x</title></head><body><pre>\n"
            "[[Page 12345]]\n"
            'Contact us at <a href="/cdn-cgi/l/email-protection#x">'
            "someone</a> regarding CMS &amp; PFS.\n"
            "A finding was significant (p<0.05 and >2 cm).\n"
            "</pre></body></html>"
        )
        store.attach_file(rule_key, txt)

        docs = list(iter_rule_docs(store))

        assert len(docs) == 1
        text = docs[0].text
        assert "<a href=" not in text
        assert "</a>" not in text
        assert "someone" in text
        assert "CMS & PFS" in text
        assert "[[Page 12345]]" in text
        assert "p<0.05 and >2 cm" in text

    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

        assert _default_root() == ROOT / ".state" / "comments"


def test_dated_versions_index_only_the_newest(tmp_path):
    """IOM revisions are hoarded as <stem>.<YYYY-MM-DD><ext> beside the
    original; the corpus reads one file per version family — the newest
    — so a chapter is not indexed once per revision."""
    from bib.item import Manual
    from bib.store import Store
    from llm.source import _attachment_paths, _attachment_paths_by_key

    store = Store(tmp_path / "b.sqlite", storage_dir=tmp_path / "st")
    key = store.create(Manual(title="Ch1", url="https://x/clm104c01.pdf"))
    for name, body in (
        ("clm104c01.pdf", b"v1"),
        ("clm104c01.2026-11-11.pdf", b"v2"),
        ("appendix.pdf", b"a"),
    ):
        f = tmp_path / name
        f.write_bytes(body)
        store.attach_file(key, f, title=name)
    names = [p.name for p in _attachment_paths(store, key)]
    assert names == ["clm104c01.2026-11-11.pdf", "appendix.pdf"]
    assert [p.name for p in _attachment_paths_by_key(store)[key]] == names