stack/tests/llm/test_iomlink.py

"""llm.iomlink — cited IOM sections as companion sources."""

from __future__ import annotations

from unittest.mock import MagicMock

from llm.iomlink import companions, lookup
from llm.links import as_source, iom_refs_of

RULE_TEXT = (
    "We refer readers to Pub. 100-04, Chapter 12, Section 30.6.4 for incident-to "
    "billing, and to the Medicare Benefit Policy Manual, Chapter 15, for coverage. "
    "See also Pub. 100-04, Chapter 12, Section 30.6.4 again."
)


class TestIomRefsOnSources:
    def test_rule_source_carries_deduplicated_refs(self):
        refs = iom_refs_of(RULE_TEXT)
        assert refs == [
            {
                "pub": "100-04",
                "chapter": "12",
                "section": "30.6.4",
                "label": "IOM 100-04 ch.12 §30.6.4",
            },
            {
                "pub": "100-02",
                "chapter": "15",
                "section": "",
                "label": "IOM 100-02 ch.15",
            },
        ]
        s = as_source({"kind": "rule", "item_key": "R", "title": "t"}, RULE_TEXT, 0.4)
        assert [r["label"] for r in s["iom_refs"]] == [
            "IOM 100-04 ch.12 §30.6.4",
            "IOM 100-02 ch.15",
        ]
        assert (
            as_source({"kind": "comment", "comment_id": "C-1"}, RULE_TEXT, 0.4)[
                "iom_refs"
            ]
            == []
        )


def _engine(rows_by_sql):
    """rows_by_sql: substring of SQL → row (document, cmetadata) or None."""
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    calls = []

    def execute(sql, params=None):
        calls.append((str(sql), params))
        r = MagicMock()
        for needle, row in rows_by_sql.items():
            if needle in str(sql):
                r.fetchone.return_value = row
                return r
        r.fetchone.return_value = None
        return r

    conn.execute.side_effect = execute
    return engine, conn, calls


MD = {
    "doctype": "manual",
    "kind": "corpus",
    "manual": "100-04",
    "chapter": "12",
    "iom_section": "30.6.4",
    "iom_title": "E/M Services Furnished Incident to",
    "page": "39",
    "item_key": "SPJB4FX2",
    "seq": "73",
    "url": "https://www.cms.gov/x/clm104c12.pdf",
    "title": "Medicare Claims Processing Manual — Chapter 12",
}


class TestCompanions:
    def _rule(self, label, text=RULE_TEXT):
        return as_source(
            {
                "kind": "rule",
                "item_key": "R",
                "p_id": "5",
                "page": "43949",
                "ordinal": "4",
                "fr_volume": "91",
                "title": "t",
                "html_url": "https://fr/x",
            },
            text,
            0.5,
        ) | {"label": label}

    def test_cited_section_becomes_a_companion_with_cited_by(self):
        engine, conn, calls = _engine(
            {"iom_section' = :section": ("## 30.6.4 - E/M ... incident to", MD)}
        )
        out = companions(
            engine,
            [self._rule("91 FR 43949 ¶4"), self._rule("91 FR 43950 ¶2")],
            max_total=4,
        )
        assert len(out) == 1
        c = out[0]
        assert c["label"] == "IOM 100-04 ch.12 §30.6.4" and c["companion"] == "iom"
        assert c["cited_by"] == ["91 FR 43949 ¶4", "91 FR 43950 ¶2"]
        assert (
            c["url"] == "https://www.cms.gov/x/clm104c12.pdf#page=39"
            and c["doctype"] == "manual"
        )
        assert c["snippet"].startswith("## 30.6.4")

    def test_prefix_and_chapter_fallbacks(self):
        finer = dict(MD, iom_section="30.6.4.1")
        engine, conn, calls = _engine({"LIKE :prefix": ("finer text", finer)})
        assert lookup(conn, "100-04", "12", "30.6.4") == ("finer text", finer)
        assert [p["prefix"] for _, p in calls if p and "prefix" in p] == ["30.6.4.%"]
        first = dict(MD, iom_section="10")
        engine, conn, calls = _engine({"<> ''": ("chapter first", first)})
        assert lookup(conn, "100-02", "15", "") == ("chapter first", first)
        engine, conn, calls = _engine({})
        assert lookup(conn, "100-04", "99", "1") is None

    def test_ranking_dedupe_cap_and_present_sections(self):
        by = {
            "iom_section' = :section": ("sec text", MD),
            "<> ''": (
                "chapter text",
                dict(MD, chapter="15", manual="100-02", iom_section="10"),
            ),
        }
        engine, conn, calls = _engine(by)
        rules = [self._rule("A"), self._rule("B")]
        out = companions(engine, rules, max_total=1)
        assert [c["label"] for c in out] == [
            "IOM 100-04 ch.12 §30.6.4"
        ]  # cited twice, sectioned → first; cap 1
        # already among the sources → not repeated
        present = as_source({k: str(v) for k, v in MD.items()}, "sec text", 0.9)
        assert companions(engine, rules + [present], max_total=4) == [
            c
            for c in companions(engine, rules + [present], max_total=4)
            if c["label"] != "IOM 100-04 ch.12 §30.6.4"
        ]

    def test_no_rule_sources_no_queries_and_failures_swallowed(self):
        engine, conn, calls = _engine({})
        assert (
            companions(
                engine,
                [
                    as_source(
                        {"kind": "comment", "comment_id": "C"},
                        "Pub. 100-04, Chapter 12",
                        0.1,
                    )
                ],
            )
            == []
        )
        assert calls == []
        engine = MagicMock()
        engine.begin.side_effect = RuntimeError("db down")
        assert companions(engine, [self._rule("A")]) == []