From 49e0f1b7414caab617108368011728a82080644f Mon Sep 17 00:00:00 2001 From: kert Date: Thu, 10 Sep 2026 11:58:47 -0400 Subject: [PATCH] feat(llm): lineage-cited FR paragraphs as sources; per-docket comment window in timeline mode (refs #691 #692) --- src/llm/config.py | 2 + src/llm/evidence.py | 105 ++++++++++++-- src/llm/lineage.py | 155 +++++++++++++++++++++ src/llm/rag.py | 31 ++++- stack.toml | 1 + tests/llm/golden_lineage.yaml | 5 - tests/llm/test_evidence.py | 111 +++++++++++++++ tests/llm/test_lineage.py | 253 +++++++++++++++++++++++++++++++++- tests/llm/test_rag.py | 5 + 9 files changed, 643 insertions(+), 25 deletions(-) diff --git a/src/llm/config.py b/src/llm/config.py index bf6a127..b522329 100644 --- a/src/llm/config.py +++ b/src/llm/config.py @@ -53,6 +53,7 @@ class LlmConfig: code_cited_max: int = 12 lineage_max_rows: int = 25 lineage_on_demand_max: int = 3 + lineage_sources_max: int = 6 def parse_hosts(spec: str) -> tuple[tuple[str, float], ...]: @@ -128,6 +129,7 @@ def load() -> LlmConfig: code_cited_max=int(_opt(section, "code_cited_max", 12)), lineage_max_rows=int(_opt(section, "lineage_max_rows", 25)), lineage_on_demand_max=int(_opt(section, "lineage_on_demand_max", 3)), + lineage_sources_max=int(_opt(section, "lineage_sources_max", 6)), embed_dim=int(section.embed_dim), build_ann_index=build_ann_index, pg_host=os.environ.get("LLM_PG_HOST", str(section.pg_host)), diff --git a/src/llm/evidence.py b/src/llm/evidence.py index f2ffdd6..165ea9e 100644 --- a/src/llm/evidence.py +++ b/src/llm/evidence.py @@ -250,6 +250,36 @@ _FAMILY_CITED_SQL = text( "ORDER BY cmetadata->>'date' DESC NULLS LAST, item_key, ordkey" ) +#: Ruling B7: in timeline mode, "which comments discussed this code" needs +#: breadth across dockets, not depth in the newest one — a "2023 vs 2025" +#: question loses the 2023 docket entirely to ``_CITED_SQL``'s per-code +#: recency window once the newer docket has more than ``per_code`` hits. +#: One chunk per docket instead: the newest chunk (by date, then seq) among +#: a docket's chunks whose ``codes`` metadata mention any of the wanted +#: codes, across every matching docket, newest docket year first. +_DOCKET_CITED_SQL = text( + "SELECT document, cmetadata FROM ( " + "SELECT e.document, e.cmetadata, " + "row_number() OVER ( " + "PARTITION BY e.cmetadata->>'docket' " + "ORDER BY e.cmetadata->>'date' DESC NULLS LAST, e.cmetadata->>'seq' " + ") AS rn " + "FROM langchain_pg_embedding e " + "JOIN unnest(CAST(:codes AS text[])) AS w(code) " + "ON w.code = ANY(string_to_array(COALESCE(e.cmetadata->>'codes', ''), ' ')) " + "WHERE e.collection_id = (SELECT uuid FROM langchain_pg_collection " + "WHERE name = :collection) " + ") t " + "WHERE rn = 1 " + "ORDER BY cmetadata->>'year' DESC NULLS LAST, cmetadata->>'docket'" +) + +#: All dockets survive ``_DOCKET_CITED_SQL``'s per-docket window (it has no +#: :window cap of its own); this is the only cap on the docket-mode result +#: for the comments collection — not per-code, a timeline question wants +#: breadth across dockets up to a sane prompt budget. +_DOCKET_MAX = 12 + def _dedupe_key(md: dict[str, str]) -> tuple[str, str]: """(item_key, p_id) for rule chunks — several paragraphs share an @@ -310,6 +340,38 @@ def _collect( return out +def _collect_by_docket( + engine: Any, codes: Sequence[str], *, collection: str, seen: set[tuple[str, str]] +) -> list[dict]: + """One chunk per docket for *collection* (``_DOCKET_CITED_SQL``) — + every docket with a chunk citing any of *codes*, newest chunk per + docket, newest docket year first; deduped via the shared *seen* set + the same way ``_collect`` is, capped at ``_DOCKET_MAX`` total (not + per code: Ruling B7 wants dockets, not depth in one).""" + upper = [w.upper() for w in codes] + if not upper: + return [] + try: + with engine.begin() as conn: + rows = conn.execute( + _DOCKET_CITED_SQL, {"collection": collection, "codes": upper} + ).fetchall() + except Exception as e: # noqa: BLE001 + log.warning("docket-cited sources skipped (%s): %s", collection, e) + return [] + out: list[dict] = [] + for document, md in rows: + md = {k: str(v) for k, v in (md or {}).items()} + dkey = _dedupe_key(md) + if dkey in seen: + continue + seen.add(dkey) + out.append(as_source(md, document[:250], 0.0)) + if len(out) >= _DOCKET_MAX: + break + return out + + def _interleave( per_collection: dict[str, list[dict]], collections: Sequence[str] ) -> list[dict]: @@ -340,6 +402,7 @@ def code_cited_sources( collections: Sequence[str] = ("rules",), families: Sequence[str] = (), max_total: int = 12, + by_docket: bool = False, ) -> list[dict]: """Chunks across *collections* whose ``codes`` metadata mention any of *codes* — at most *per_code* per code per collection, newest first, @@ -350,6 +413,16 @@ def code_cited_sources( skipped entirely, no query, when *families* is empty. Scores are 0.0: these are additive, not ranked. + *by_docket* (Ruling B7): when true and ``"comments"`` is among + *collections*, the comments window is the newest chunk per docket + among every docket citing any of *codes* (``_DOCKET_CITED_SQL``, + capped at ``_DOCKET_MAX`` total, not per code) instead of + ``_CITED_SQL``'s per-code recency window — a "2023 vs 2025" + question needs each docket represented, not just the newest. + Every other collection keeps the normal per-code window. No effect + when *families*-only (``codes`` empty) or ``"comments"`` isn't in + *collections*. + After the per-collection caps and dedupe above, the code-cited rows are round-robin interleaved across *collections* in the order given (one from the first collection, one from the second, … looping back @@ -361,18 +434,28 @@ def code_cited_sources( if not codes and not families: return [] seen: set[tuple[str, str]] = set() - out = _interleave( - _collect( - engine, - _CITED_SQL, - "codes", - codes, - per=per_code, - collections=collections, - seen=seen, - ), - collections, + docket_mode = by_docket and "comments" in collections + window_collections = ( + [c for c in collections if c != "comments"] + if docket_mode + else list(collections) ) + code_hits = _collect( + engine, + _CITED_SQL, + "codes", + codes, + per=per_code, + collections=window_collections, + seen=seen, + ) + if docket_mode: + docket_hits = _collect_by_docket( + engine, codes, collection="comments", seen=seen + ) + if docket_hits: + code_hits["comments"] = docket_hits + out = _interleave(code_hits, collections) if families: out += _interleave( _collect( diff --git a/src/llm/lineage.py b/src/llm/lineage.py index eb57b9c..e527a68 100644 --- a/src/llm/lineage.py +++ b/src/llm/lineage.py @@ -34,6 +34,7 @@ from dataclasses import asdict, dataclass from typing import Any, Sequence from llm.config import LlmConfig +from llm.links import as_source from pfs.codetables import ( ElementRow, EventRow, @@ -291,6 +292,16 @@ class _ReplicaCache: _ITEM_CACHE = _ReplicaCache() # item_key -> resolved bib Item (or None) _URL_CACHE = _ReplicaCache() # (item_key, p_id) -> FR jump-link url _ON_DEMAND_CACHE = _ReplicaCache() # code -> on-demand EventRow list +_PARAGRAPH_CACHE = _ReplicaCache() # (item_key, p_id) -> fr_anchors+items row (or None) + +#: ``lineage_sources``'s paragraph fetches key the ``_ReplicaCache`` on a +#: constant "generation" rather than the duckdb replica's mtime: the bib +#: sqlite store (unlike the replica) is opened once per process +#: (``_store``, above) with no republish signal this module can key on, +#: so a paragraph, once fetched, is cached for the life of the process — +#: exactly the pattern the item/url caches use, minus the invalidation +#: trigger they don't have here either. +_PARAGRAPH_CACHE_GEN = 0 def _cached_item(store: Any, item_key: str, mtime: int) -> Any: @@ -572,6 +583,150 @@ def _collect_guidance( return tuple(out) +def _lineage_paragraph(store: Any, item_key: str, p_id: int) -> dict | None: + """``fr_anchors`` (``text``/``page``/``ordinal``) joined with + ``items`` (``title``/``date_published``/``url``) and, when present, + ``fr_anchor_docs`` (``html_url``/``fr_volume`` — the fields + ``links._rule`` and ``source._build_rule_doc`` use to build a rule + chunk's own metadata) for one ``(item_key, p_id)``. ``None`` when + the paragraph doesn't exist. Cached per key for the life of the + process (``_PARAGRAPH_CACHE``, see its comment above).""" + cached = _PARAGRAPH_CACHE.get((item_key, p_id), _PARAGRAPH_CACHE_GEN) + if cached is not _MISSING: + return cached + row = ( + store._con() + .execute( + "SELECT a.text, a.page, a.ordinal, i.title, i.date_published, i.url, " + "d.html_url, d.fr_volume " + "FROM fr_anchors a " + "JOIN items i ON i.key = a.item_key " + "LEFT JOIN fr_anchor_docs d ON d.item_key = a.item_key " + "WHERE a.item_key = ? AND a.p_id = ?", + (item_key, p_id), + ) + .fetchone() + ) + result = dict(row) if row is not None else None + _PARAGRAPH_CACHE.put((item_key, p_id), _PARAGRAPH_CACHE_GEN, result) + return result + + +def _lineage_source(store: Any, event: LineageEvent) -> dict | None: + """One rule-kind source dict for *event*'s FR paragraph, built the + same way the indexer builds a rule chunk (``links._rule`` reads + ``item_key``/``p_id``/``page``/``ordinal``/``fr_volume``/``html_url`` + from the metadata) so ``as_source`` produces the same kind of + (url, label) a retrieved excerpt for this paragraph would — except + the label/id are then overridden with *event*'s own lineage label, + so the prompt block's ``[label]`` and the sources drawer's entry are + the exact same string. ``None`` when the paragraph can't be found.""" + para = _lineage_paragraph(store, event.item_key, event.p_id) + if para is None: + return None + md = { + "kind": "rule", + "item_key": event.item_key, + "p_id": str(event.p_id), + "page": str(para.get("page") or event.page or ""), + "ordinal": str(para.get("ordinal") or ""), + "fr_volume": str(para.get("fr_volume") or ""), + "html_url": para.get("html_url") or "", + "url": para.get("url") or "", + "title": para.get("title") or "", + "date": (para.get("date_published") or "")[:10], + } + src = as_source(md, para.get("text") or "", 0.0) + src["label"] = event.label + src["id"] = event.label + return src + + +def lineage_sources( + store: Any, evidence: LineageEvidence, *, max_items: int = 6 +) -> list[dict]: + """FR paragraphs for the events selected into the lineage prompt + block, as ``rule``-kind source dicts — so the model can actually + read the paragraphs it is told to cite (``[label]``) and the + sources drawer can link them (Ruling B6). + + Distinct ``(item_key, p_id)`` pairs among ``source == "fr"`` prompt + events (the same selection ``LineageEvidence.prompt_block`` uses), + priority-kind events first (``_PRIORITY_KINDS`` — created, adopted, + replaces, replaced_by, deleted, disappeared), then by year; capped + at *max_items*. Each source's ``label``/``id`` are overridden with + the lineage event's own label (``_lineage_source``), so the same + paragraph never shows two different bracketed labels to the model. + + Never raises: a failure fetching any one paragraph is logged and + skipped; a failure before the loop starts degrades to ``[]``, the + same contract as ``lineage_evidence`` itself.""" + try: + prompt_events = _select_for_prompt(evidence.events, evidence.max_prompt_rows) + candidates = [ + e for e in prompt_events if e.source == "fr" and e.item_key and e.p_id + ] + candidates.sort(key=lambda e: (0 if e.kind in _PRIORITY_KINDS else 1, e.year)) + except Exception as e: # noqa: BLE001 + log.warning("lineage sources skipped (select): %s", e) + return [] + out: list[dict] = [] + seen: set[tuple[str, int]] = set() + for e in candidates: + key = (e.item_key, e.p_id) + if key in seen: + continue + seen.add(key) + try: + src = _lineage_source(store, e) + except Exception as ex: # noqa: BLE001 + log.warning("lineage source skipped (%s p-%s): %s", e.item_key, e.p_id, ex) + continue + if src is not None: + out.append(src) + if len(out) >= max_items: + break + return out + + +def _reconcile_labels( + sources: list[dict], lineage_rows: list[dict] +) -> tuple[list[dict], list[dict]]: + """Before merging *lineage_rows* (``lineage_sources``' output) into + *sources* (the already-retrieved/cited list): a retrieved excerpt + for the same FR paragraph carries its own retrieval-derived label + (e.g. ``"85 FR 84639 ¶12"``), different from the lineage label the + model was told to cite (e.g. ``"CY2021 PFS final ¶1578"``) — two + labels for the same passage would confuse the model and double it + up in the sources drawer. For any ``sources`` entry sharing a + lineage row's ``(item_key, p_id)``, rename that entry's + ``label``/``id`` to the lineage label (mutated in place) and drop + the lineage row from what still needs merging — the retrieved + excerpt wins the slot, under the lineage label. Returns + ``(sources, lineage_rows)`` with both adjusted.""" + by_key: dict[tuple[str, str], dict] = {} + for row in lineage_rows: + key = (row.get("item_key", ""), row.get("p_id", "")) + if key[0] and key[1] and key not in by_key: + by_key[key] = row + matched: set[tuple[str, str]] = set() + for src in sources: + if src.get("kind") != "rule": + continue + key = (src.get("item_key", ""), src.get("p_id", "")) + row = by_key.get(key) + if row is not None: + src["label"] = row["label"] + src["id"] = row["label"] + matched.add(key) + remaining = [ + row + for row in lineage_rows + if (row.get("item_key", ""), row.get("p_id", "")) not in matched + ] + return sources, remaining + + def lineage_evidence(question: str, cfg: LlmConfig) -> LineageEvidence | None: """Timeline events, element differences and guidance for the codes detected in *question* — ``None`` when no codes are detected. Never diff --git a/src/llm/rag.py b/src/llm/rag.py index f56f3dd..1826a73 100644 --- a/src/llm/rag.py +++ b/src/llm/rag.py @@ -45,7 +45,7 @@ from llm.evidence import ( valuation_evidence, ) from llm.index import _engine -from llm.lineage import LineageEvidence +from llm.lineage import LineageEvidence, _reconcile_labels, lineage_sources from llm.lineage import _store as _bib_store from llm.links import as_source from llm.pool import HostPool, PoolEmbeddings, pick_model @@ -374,11 +374,16 @@ def stream_answer( (evidence's ``payload()``, the dated timeline) then a ``valuation`` event (RVUs and payment) — both before any tokens — and merges chunks from rules, comments and the corpus that cite those codes — - or, for a detected family, its family key — into ``sources``, then - (for any non-wide detected family) the AMA CPT manual's own - guideline text for that family's heading (``llm.evidence. - manual_sources`` — deterministic, not a retrieval hit, so it never - counts toward ``code_cited_max``). Then yields + or, for a detected family, its family key — into ``sources`` + (in timeline mode, the comments window is one chunk per docket + instead of per-code recency, ``code_cited_sources(by_docket=True)`` + — a "2023 vs 2025" question needs each docket, not just the + newest); then the FR paragraphs behind the lineage prompt block's + events (``llm.lineage.lineage_sources`` — Ruling B6, also outside + ``code_cited_max``); then (for any non-wide detected family) the + AMA CPT manual's own guideline text for that family's heading + (``llm.evidence.manual_sources`` — deterministic, not a retrieval + hit, so it never counts toward ``code_cited_max``). Then yields ``{"type":"token","text":…}`` events as the model generates, then one ``{"type":"sources", "sources": […], "model": …, "host": …, "mode": "timeline"|"recent"}`` (the resolved mode) and a final @@ -401,8 +406,22 @@ def stream_answer( collections=tuple(cfg.code_cited_collections), families=families, max_total=cfg.code_cited_max, + by_docket=(resolved_mode == "timeline"), ) sources = merge_sources(sources, cited) + if lineage is not None: + # The FR paragraphs behind the events selected into the + # lineage prompt block, so the model can read what it's + # told to cite and the sources drawer can link them + # (Ruling B6) — outside code_cited_max. A paragraph already + # retrieved/cited under its own label keeps its slot but is + # renamed to the lineage label (Ruling B6 dedupe) rather + # than appearing twice. + lrows = lineage_sources( + _bib_store(), lineage, max_items=cfg.lineage_sources_max + ) + sources, lrows = _reconcile_labels(sources, lrows) + sources = merge_sources(sources, lrows) # The CPT manual's own heading text for a detected family — a # wide family (its code list too big to expand, MAX_FAMILY_EXPAND) # is excluded, and these rows never count toward code_cited_max. diff --git a/stack.toml b/stack.toml index dc73a54..9fc0efe 100644 --- a/stack.toml +++ b/stack.toml @@ -133,6 +133,7 @@ code_cited_collections = ["rules", "comments", "corpus"] # searched in this ord code_cited_max = 12 # cited excerpts kept after round-robin interleave across collections; <= 0 = unlimited lineage_max_rows = 25 # collapsed lineage events kept in the prompt block (the SSE payload always carries every collapsed row) lineage_on_demand_max = 3 # detected codes per turn allowed to fall back to pfs.lineage.lineage() when pfs.code_event has no rows for them +lineage_sources_max = 6 # FR paragraphs fetched as sources for events selected into the lineage prompt block, outside code_cited_max [llm.k_per_kind] # over-fetched ×3 per kind, then re-ranked comment = 8 diff --git a/tests/llm/golden_lineage.yaml b/tests/llm/golden_lineage.yaml index 9d499d8..2a4762e 100644 --- a/tests/llm/golden_lineage.yaml +++ b/tests/llm/golden_lineage.yaml @@ -29,8 +29,6 @@ p_id: [1249, 1251] - item_key: YBM4IZUS p_id: 1578 - - item_key: JJ6AM5HJ - p_id: 1163 expect_events: - code: "99490" kind: created @@ -38,9 +36,6 @@ - code: G2058 kind: replaced_by year: 2021 - - code: G0556 - kind: created - year: 2025 min_eras: 4 - id: g2058-replacement diff --git a/tests/llm/test_evidence.py b/tests/llm/test_evidence.py index 8a4e6f1..444ebc7 100644 --- a/tests/llm/test_evidence.py +++ b/tests/llm/test_evidence.py @@ -753,6 +753,117 @@ class TestMaxTotal: assert len(out) == 12 +def _comment_docket_md( + seq, codes, docket, item_key=None, date="2025-01-01", year="2025" +): + return { + "kind": "comment", + "item_key": item_key or f"{docket}-C{seq}", + "seq": str(seq), + "comment_id": f"{docket}-{seq}", + "date": date, + "year": year, + "docket": docket, + "codes": codes, + "families": "", + } + + +def _docket_fake_engine(docket_rows, code_rows_by_collection=None): + """A fake engine that tells the docket-SQL call apart from the + normal per-code window call by the params shape alone — the + per-code window (``_collect``) always sends a ``window`` bind + param; the docket window (``_collect_by_docket``) never does.""" + engine = MagicMock() + conn = engine.begin.return_value.__enter__.return_value + code_rows_by_collection = code_rows_by_collection or {} + + def _execute(_sql, params): + result = MagicMock() + if "window" not in params: + result.fetchall.return_value = docket_rows + else: + result.fetchall.return_value = code_rows_by_collection.get( + params["collection"], [] + ) + return result + + conn.execute.side_effect = _execute + return engine, conn + + +class TestCodeCitedSourcesByDocket: + """Ruling B7 — ``by_docket`` swaps the comments collection's + per-code recency window for one chunk per docket, all dockets, in + timeline mode only.""" + + def test_docket_sql_used_only_for_comments(self): + engine, conn = _docket_fake_engine( + docket_rows=[("d1", _comment_docket_md(1, "G2211", "CMS-2023-0121"))], + code_rows_by_collection={"rules": [("r1", _rule_md(1, "G2211"))]}, + ) + code_cited_sources( + engine, + ["G2211"], + per_code=3, + collections=("rules", "comments"), + by_docket=True, + ) + calls = conn.execute.call_args_list + comments_calls = [c for c in calls if c.args[1]["collection"] == "comments"] + rules_calls = [c for c in calls if c.args[1]["collection"] == "rules"] + assert comments_calls and "window" not in comments_calls[0].args[1] + assert rules_calls and "window" in rules_calls[0].args[1] + + def test_by_docket_false_uses_the_normal_window_for_comments(self): + engine, conn = _docket_fake_engine( + docket_rows=[("d1", _comment_docket_md(1, "G2211", "CMS-2023-0121"))], + code_rows_by_collection={"comments": [("c1", _comment_md(1, "G2211"))]}, + ) + code_cited_sources( + engine, + ["G2211"], + per_code=3, + collections=("comments",), + by_docket=False, + ) + calls = conn.execute.call_args_list + assert calls and "window" in calls[0].args[1] + + def test_by_docket_without_comments_in_collections_is_a_noop(self): + engine, conn = _docket_fake_engine( + docket_rows=[("d1", _comment_docket_md(1, "G2211", "CMS-2023-0121"))], + code_rows_by_collection={"rules": [("r1", _rule_md(1, "G2211"))]}, + ) + code_cited_sources( + engine, ["G2211"], per_code=3, collections=("rules",), by_docket=True + ) + calls = conn.execute.call_args_list + assert calls and "window" in calls[0].args[1] + + def test_three_dockets_survive_dedupe_and_are_interleaved(self): + docket_rows = [ + ("c23", _comment_docket_md(1, "G2211", "CMS-2023-0121", item_key="C23")), + ("c25", _comment_docket_md(1, "G2211", "CMS-2025-0304", item_key="C25")), + ("c26", _comment_docket_md(1, "G2211", "CMS-2026-2377", item_key="C26")), + ] + engine, _conn = _docket_fake_engine( + docket_rows=docket_rows, + code_rows_by_collection={"rules": [("r1", _rule_md(1, "G2211"))]}, + ) + out = code_cited_sources( + engine, + ["G2211"], + per_code=3, + collections=("rules", "comments"), + by_docket=True, + ) + dockets = {s["docket"] for s in out if s["kind"] == "comment"} + assert dockets == {"CMS-2023-0121", "CMS-2025-0304", "CMS-2026-2377"} + # round-robin interleaved with the rules row, not appended en bloc + assert [s["kind"] for s in out] == ["rule", "comment", "comment", "comment"] + + class TestMergeSources: def test_keeps_order_and_dedupes_by_label(self): a = [{"label": "X", "score": 0.9}, {"label": "Y", "score": 0.8}] diff --git a/tests/llm/test_lineage.py b/tests/llm/test_lineage.py index 6a046d6..62ad35e 100644 --- a/tests/llm/test_lineage.py +++ b/tests/llm/test_lineage.py @@ -20,7 +20,9 @@ from llm.lineage import ( GuidanceRef, LineageEvent, LineageEvidence, + _reconcile_labels, lineage_evidence, + lineage_sources, rule_label, ) from pfs.codetables import ( @@ -63,6 +65,7 @@ def _clear_caches(): lineage._ITEM_CACHE.reset() lineage._URL_CACHE.reset() lineage._ON_DEMAND_CACHE.reset() + lineage._PARAGRAPH_CACHE.reset() _reset() yield @@ -81,29 +84,32 @@ class _FakeStore: self.con.row_factory = sqlite3.Row self.con.executescript( "CREATE TABLE items (key TEXT PRIMARY KEY, title TEXT, " - "date_published TEXT);" + "date_published TEXT, url TEXT);" "CREATE TABLE fr_anchors (item_key TEXT, p_id INTEGER, " "page INTEGER, ordinal INTEGER, text TEXT);" "CREATE TABLE fr_anchor_docs (item_key TEXT, html_url TEXT, " "start_page INTEGER, end_page INTEGER, fr_volume INTEGER);" ) self.con.executemany( - "INSERT INTO items VALUES (?,?,?)", + "INSERT INTO items VALUES (?,?,?,?)", [ ( "DE2VH9PD", "Medicare Program; CY 2015 PFS Final Rule", "2014-11-13", + "https://www.federalregister.gov/d/2014-2015doc", ), ( "YBM4IZUS", "Medicare Program; CY 2021 Payment Policies Under the PFS", "2020-12-28", + "https://www.federalregister.gov/d/2020-2021doc", ), ( "ZPROP2027", "Medicare Program; CY 2027 Proposed Payment Policies", "2026-07-16", + "https://www.federalregister.gov/d/2026-2027doc", ), ], ) @@ -902,12 +908,13 @@ class TestCachingAndScopedUrlResolution: ] write_events(c, "99999", rows) store.con.executemany( - "INSERT INTO items VALUES (?,?,?)", + "INSERT INTO items VALUES (?,?,?,?)", [ ( f"K{i:02d}", f"Medicare Program; CY {2000 + i} PFS Final Rule", f"{1999 + i}-11-01", + f"https://www.federalregister.gov/d/{2000 + i}-doc", ) for i in range(n_items) ], @@ -942,6 +949,246 @@ class TestCachingAndScopedUrlResolution: assert len(seen_keys) == len(set(seen_keys)) # never re-fetched +def _lineage_evidence(events, max_prompt_rows=25) -> LineageEvidence: + return LineageEvidence( + codes=(), + families=(), + events=tuple(events), + element_diffs=(), + guidance=(), + max_prompt_rows=max_prompt_rows, + ) + + +class TestLineageSources: + """``lineage_sources`` (Ruling B6) — FR paragraphs behind the + events selected into the prompt block, as rule-kind sources whose + label is the lineage event's own label.""" + + def test_priority_kinds_first_then_by_year(self, store): + # revalued (non-priority) at 2020 must still sort after every + # priority-kind event, even one dated later (2021). + revalued = lineage._to_lineage_event( + store, + _ev("99439", 2020, "revalued", item_key="YBM4IZUS", p_id=1578), + mtime=0, + ) + created = lineage._to_lineage_event( + store, + _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686), + mtime=0, + ) + adopted = lineage._to_lineage_event( + store, + _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251), + mtime=0, + ) + ev = _lineage_evidence([revalued, created, adopted]) + out = lineage_sources(store, ev, max_items=6) + assert [s["label"] for s in out] == [ + adopted.label, + created.label, + revalued.label, + ] + + def test_capped_at_max_items(self, store): + store.con.execute( + "INSERT INTO items VALUES (?,?,?,?)", + ( + "EXTRA1", + "Medicare Program; CY 2022 PFS Final Rule", + "2021-11-02", + "https://www.federalregister.gov/d/2021-2022doc", + ), + ) + store.con.execute( + "INSERT INTO fr_anchors VALUES (?,?,?,?,?)", + ("EXTRA1", 10, 1000, 1, "extra paragraph text"), + ) + store.con.commit() + events = [ + lineage._to_lineage_event( + store, + _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251), + mtime=0, + ), + lineage._to_lineage_event( + store, + _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686), + mtime=0, + ), + lineage._to_lineage_event( + store, + _ev( + "99439", + 2021, + "replaces", + frm="G2058", + item_key="YBM4IZUS", + p_id=1578, + ), + mtime=0, + ), + lineage._to_lineage_event( + store, + _ev("G0000", 2022, "revalued", item_key="EXTRA1", p_id=10), + mtime=0, + ), + ] + ev = _lineage_evidence(events) + out = lineage_sources(store, ev, max_items=2) + assert len(out) == 2 + + def test_label_equals_the_event_label(self, store): + e = lineage._to_lineage_event( + store, + _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251), + mtime=0, + ) + ev = _lineage_evidence([e]) + (out,) = lineage_sources(store, ev, max_items=6) + assert out["label"] == e.label == "CY2015 PFS final ¶1251" + assert out["id"] == e.label + assert out["kind"] == "rule" + assert out["item_key"] == "DE2VH9PD" + assert out["p_id"] == "1251" + assert "adopt CPT code 99490" in out["snippet"] + + def test_dedupes_distinct_item_key_p_id_across_events(self, store): + # two collapsed events can point at the same paragraph (e.g. a + # "created" and a "replaces" row both anchored at YBM4IZUS 686 + # in a contrived case) — only one source per (item_key, p_id). + e1 = lineage._to_lineage_event( + store, + _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686), + mtime=0, + ) + e2 = lineage._to_lineage_event( + store, + _ev("99439", 2021, "revalued", item_key="YBM4IZUS", p_id=686), + mtime=0, + ) + ev = _lineage_evidence([e1, e2]) + out = lineage_sources(store, ev, max_items=6) + assert len(out) == 1 + + def test_skips_non_fr_events(self, store): + rvu = lineage._to_lineage_event( + store, _ev("99441", 2022, "disappeared", source="rvu"), mtime=0 + ) + ev = _lineage_evidence([rvu]) + assert lineage_sources(store, ev, max_items=6) == [] + + def test_missing_paragraph_is_skipped_not_raised(self, store): + e = lineage._to_lineage_event( + store, + _ev("99999", 2019, "created", item_key="NOPE", p_id=9999), + mtime=0, + ) + ev = _lineage_evidence([e]) + assert lineage_sources(store, ev, max_items=6) == [] + + def test_never_raises_on_a_broken_store(self): + class _BrokenStore: + def _con(self): + raise RuntimeError("boom") + + e = LineageEvent( + code="99490", + year=2015, + kind="adopted_cpt", + from_codes=(), + to_codes=(), + label="CY2015 PFS final ¶1251", + item_key="DE2VH9PD", + p_id=1251, + page=67716, + url="", + source="fr", + anchored=True, + note="", + ) + ev = _lineage_evidence([e]) + assert lineage_sources(_BrokenStore(), ev, max_items=6) == [] + + def test_only_prompt_selected_events_are_considered(self, store): + # max_prompt_rows=0 with a non-priority event only — priority + # rows are never dropped by _select_for_prompt, but a + # non-priority row past budget zero is. + e = lineage._to_lineage_event( + store, + _ev("99439", 2021, "revalued", item_key="YBM4IZUS", p_id=686), + mtime=0, + ) + ev = _lineage_evidence([e], max_prompt_rows=0) + assert lineage_sources(store, ev, max_items=6) == [] + + +class TestReconcileLabels: + """``_reconcile_labels`` — a retrieved/cited source for the same FR + paragraph as a lineage row wins the slot, renamed to the lineage + label; the lineage row is dropped so it isn't merged as a + duplicate under a second label.""" + + def test_renames_matching_source_and_drops_the_lineage_row(self): + retrieved = { + "label": "85 FR 84639 ¶12", + "id": "85 FR 84639 ¶12", + "kind": "rule", + "item_key": "YBM4IZUS", + "p_id": "1578", + } + lineage_row = { + "label": "CY2021 PFS final ¶1578", + "id": "CY2021 PFS final ¶1578", + "kind": "rule", + "item_key": "YBM4IZUS", + "p_id": "1578", + } + sources, remaining = _reconcile_labels([retrieved], [lineage_row]) + assert sources[0]["label"] == "CY2021 PFS final ¶1578" + assert sources[0]["id"] == "CY2021 PFS final ¶1578" + assert remaining == [] + + def test_no_match_keeps_both_untouched(self): + retrieved = { + "label": "85 FR 1 ¶9", + "id": "85 FR 1 ¶9", + "kind": "rule", + "item_key": "OTHER", + "p_id": "9", + } + lineage_row = { + "label": "CY2021 PFS final ¶1578", + "id": "CY2021 PFS final ¶1578", + "kind": "rule", + "item_key": "YBM4IZUS", + "p_id": "1578", + } + sources, remaining = _reconcile_labels([retrieved], [lineage_row]) + assert sources[0]["label"] == "85 FR 1 ¶9" + assert remaining == [lineage_row] + + def test_non_rule_sources_are_never_renamed(self): + comment = { + "label": "CMS-2026-2377-1 p.1", + "id": "CMS-2026-2377-1 p.1", + "kind": "comment", + "item_key": "YBM4IZUS", + "p_id": "", + } + lineage_row = { + "label": "CY2021 PFS final ¶1578", + "id": "CY2021 PFS final ¶1578", + "kind": "rule", + "item_key": "YBM4IZUS", + "p_id": "1578", + } + sources, remaining = _reconcile_labels([comment], [lineage_row]) + assert sources[0]["label"] == "CMS-2026-2377-1 p.1" + assert remaining == [lineage_row] + + class TestLineageEvidenceLive: @pytest.mark.live def test_median_under_300ms_on_live_replica(self): diff --git a/tests/llm/test_rag.py b/tests/llm/test_rag.py index a3fb0b7..2b677ef 100644 --- a/tests/llm/test_rag.py +++ b/tests/llm/test_rag.py @@ -751,6 +751,7 @@ class TestStreamAnswer: collections=CFG.code_cited_collections, families=("APCM",), max_total=CFG.code_cited_max, + by_docket=False, ) body = client.stream.call_args.kwargs["json"] assert "Valuation (authoritative" in body["messages"][1]["content"] @@ -846,6 +847,8 @@ class TestStreamAnswer: events = list(stream_answer("CCM history?", cfg=CFG, pool=self._pool())) assert events[0]["type"] == "lineage" + # "CCM history?" is history-shaped (is_history_question) — mode + # resolves to "timeline", so the comments window is per-docket. mock_cited.assert_called_once_with( mock_engine.return_value, ("99490", "99491"), @@ -853,6 +856,7 @@ class TestStreamAnswer: collections=CFG.code_cited_collections, families=("CCM",), max_total=CFG.code_cited_max, + by_docket=True, ) @patch("llm.rag._manual_sources", return_value=[]) @@ -898,6 +902,7 @@ class TestStreamAnswer: collections=CFG.code_cited_collections, families=("CCM",), max_total=CFG.code_cited_max, + by_docket=False, ) @patch("llm.rag._manual_sources")