Files
stack/tests/llm/test_lineage.py

1374 lines
48 KiB
Python

"""llm.lineage — dated timeline, element diffs and guidance as a
lineage SSE event + cited prompt block."""
from __future__ import annotations
import json
import sqlite3
import time
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import patch
import duckdb
import pytest
import llm.lineage as lineage
from llm.config import LlmConfig
from llm.lineage import (
ElementDiff,
GuidanceRef,
LineageEvent,
LineageEvidence,
_reconcile_labels,
lineage_evidence,
lineage_sources,
rule_label,
)
from pfs.codetables import (
ElementRow,
EventRow,
GuidanceRow,
ensure_tables,
write_elements,
write_events,
write_guidance,
)
from pfs.families import Detection
CFG = LlmConfig(
ollama_hosts=("http://h1:11434",),
embed_model="e",
instruct_model="c",
embed_dim=768,
build_ann_index=False,
pg_host="x",
pg_port=5432,
pg_db="llm",
pg_user="llm",
duckdb_replica="/nonexistent/aco.ro.duckdb",
lineage_max_rows=25,
lineage_on_demand_max=3,
)
@pytest.fixture(autouse=True)
def _clear_caches():
"""Every process-global cache lineage.py keeps is shared across
tests (and with llm.evidence's replica cache) — no test may inherit
another's handle, store, or on-demand memo."""
import llm.evidence as evidence
def _reset():
evidence._REPLICA = None
lineage._STORE = None
lineage._ITEM_CACHE.reset()
lineage._URL_CACHE.reset()
lineage._ON_DEMAND_CACHE.reset()
lineage._PARAGRAPH_CACHE.reset()
_reset()
yield
_reset()
class _FakeStore:
"""A bib.Store stand-in: ``items``/``fr_anchors``/``fr_anchor_docs``
tables the way bib.frlink.resolve and Store.get expect, seeded with
the real CY2015/CY2021 CCM history (99490 adopted in 2015, G2058
replaced by the new CPT code 99439 in 2021) already used as a fixture
in tests/pfs/test_lineage.py."""
def __init__(self):
self.con = sqlite3.connect(":memory:")
self.con.row_factory = sqlite3.Row
self.con.executescript(
"CREATE TABLE items (key TEXT PRIMARY KEY, title 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 (?,?,?,?)",
[
(
"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",
),
],
)
self.con.executemany(
"INSERT INTO fr_anchor_docs VALUES (?,?,?,?,?)",
[
("DE2VH9PD", "https://fr.test/2015-doc", 67000, 68000, 79),
("YBM4IZUS", "https://fr.test/2021-doc", 84000, 85000, 85),
],
)
self.con.executemany(
"INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
[
(
"DE2VH9PD",
1251,
67716,
1,
"Accordingly, we will adopt CPT code 99490 for Medicare "
"CCM services.",
),
(
"YBM4IZUS",
686,
84547,
1,
"We are finalizing HCPCS code G2058 as new CPT code 99439.",
),
(
"YBM4IZUS",
1578,
84639,
2,
"A temporary crosswalk between G2058 and new CPT code "
"99439 (with a descriptor identical to G2058).",
),
],
)
def _con(self):
return self.con
def get(self, key: str) -> SimpleNamespace:
row = self.con.execute(
"SELECT title, date_published FROM items WHERE key = ?", (key,)
).fetchone()
if row is None:
raise KeyError(key)
return SimpleNamespace(title=row["title"], date_published=row["date_published"])
def close(self):
self.con.close()
@pytest.fixture
def store():
s = _FakeStore()
yield s
s.close()
@pytest.fixture
def con(tmp_path):
"""A real temp DuckDB replica (not :memory:) so lineage_evidence's
real _connect/_replica_path path is exercised end to end, the same
way TestConcurrentChats in test_evidence.py does."""
db = tmp_path / "aco.ro.duckdb"
c = duckdb.connect(str(db))
ensure_tables(c)
yield c, db
c.close()
def _ev(
code,
year,
kind,
frm="",
to="",
item_key="",
p_id=0,
page=0,
source="fr",
anchored=True,
note="",
) -> EventRow:
return EventRow(
code, year, kind, frm, to, item_key, p_id, page, source, anchored, note
)
def _el(code, year, type_, value, item_key="", p_id=0, source="fr") -> ElementRow:
return ElementRow(code, year, type_, value, "", "", item_key, p_id, 0, source)
def _gd(
family, code, kind, locator, item_key_src, p_id_src, item_key=""
) -> GuidanceRow:
return GuidanceRow(family, code, kind, locator, item_key, item_key_src, p_id_src, 0)
class TestRuleLabel:
def test_final(self):
assert (
rule_label(
"Medicare Program; CY 2021 Payment Policies Under the PFS",
"2020-12-28",
1578,
)
== "CY2021 PFS final ¶1578"
)
def test_proposed(self):
assert (
rule_label(
"Medicare Program; CY 2027 Proposed Payment Policies",
"2026-07-16",
394,
)
== "CY2027 PFS proposed ¶394"
)
class TestCollapse:
def test_prefers_anchored_fr_over_rvu_for_the_same_group(self):
rows = [
_ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686, source="fr"),
_ev("99439", 2021, "created", source="rvu", anchored=True, note="dupe"),
]
collapsed = lineage._collapse(rows)
assert len(collapsed) == 1
assert collapsed[0].source == "fr" and collapsed[0].p_id == 686
def test_tiebreaks_on_lowest_p_id(self):
rows = [
_ev(
"99439",
2021,
"replaces",
to="",
frm="G2058",
item_key="YBM4IZUS",
p_id=2999,
source="fr",
),
_ev(
"99439",
2021,
"replaces",
frm="G2058",
item_key="YBM4IZUS",
p_id=1578,
source="fr",
),
]
collapsed = lineage._collapse(rows)
assert len(collapsed) == 1 and collapsed[0].p_id == 1578
def test_sorted_by_year_code_then_kind_order(self):
rows = [
_ev("99439", 2021, "replaces", frm="G2058", p_id=1),
_ev("G2058", 2021, "replaced_by", to="99439", p_id=1),
_ev("99439", 2015, "created", p_id=2),
]
collapsed = lineage._collapse(rows)
assert [(e.year, e.code, e.kind) for e in collapsed] == [
(2015, "99439", "created"),
(2021, "99439", "replaces"),
(2021, "G2058", "replaced_by"),
]
def test_distinct_groups_are_not_merged(self):
rows = [
_ev("99439", 2021, "created", p_id=1),
_ev("99439", 2021, "replaces", frm="G2058", p_id=1),
]
assert len(lineage._collapse(rows)) == 2
def test_created_across_three_years_collapses_to_the_earliest(self):
"""Ruling B5: G2058 is mentioned as "created" by three different
rules (2021, 2022, 2023 — each rule re-establishing the G-code
after it was crosswalked away) but it was only actually created
once."""
rows = [
_ev("G2058", 2021, "created", item_key="A", p_id=100),
_ev("G2058", 2022, "created", item_key="B", p_id=200),
_ev("G2058", 2023, "created", item_key="C", p_id=300),
]
collapsed = lineage._collapse(rows)
assert [(e.year, e.code, e.kind) for e in collapsed] == [
(2021, "G2058", "created")
]
assert collapsed[0].item_key == "A"
def test_deleted_twice_keeps_only_the_first(self):
rows = [
_ev("99213", 2019, "deleted", item_key="A", p_id=1),
_ev("99213", 2024, "deleted", item_key="B", p_id=2),
]
collapsed = lineage._collapse(rows)
assert [(e.year, e.kind) for e in collapsed] == [(2019, "deleted")]
def test_earliest_only_picks_the_preferred_row_within_that_year(self):
"""Two rows tie on the earliest year — the anchored-FR/lowest-p_id
preference still applies among them."""
rows = [
_ev(
"G2058",
2021,
"created",
item_key="LATE",
p_id=999,
source="fr",
anchored=True,
),
_ev(
"G2058",
2021,
"created",
item_key="EARLY",
p_id=100,
source="fr",
anchored=True,
),
_ev("G2058", 2022, "created", item_key="OTHER", p_id=1),
]
collapsed = lineage._collapse(rows)
assert len(collapsed) == 1
assert collapsed[0].year == 2021 and collapsed[0].item_key == "EARLY"
def test_replaced_by_rows_in_two_years_both_survive(self):
"""replaced_by is directional/repeatable, not a point-in-time
kind — a code can legitimately be replaced by different things
in different years, so both rows must survive collapse."""
rows = [
_ev("99439", 2021, "replaced_by", to="G2058", p_id=1),
_ev("99439", 2022, "replaced_by", to="G2214", p_id=2),
]
collapsed = lineage._collapse(rows)
assert [(e.year, e.to_codes) for e in collapsed] == [
(2021, "G2058"),
(2022, "G2214"),
]
class TestElementDiffs:
def test_no_diffs_when_fewer_than_two_codes_have_elements(self, con):
c, _ = con
assert lineage._element_diffs(c, ["99490"], None, 1) == ((), "")
def test_notes_when_exactly_one_code_has_elements(self, con):
c, _ = con
write_elements(c, "99490", [_el("99490", 2026, "consent", "required")], [])
diffs, note = lineage._element_diffs(c, ["99490", "99491", "G0556"], None, 1)
assert diffs == ()
assert note == "elements extracted for 1 of 3 codes"
def test_shared_value_diffs_from_the_one_missing_code(self, con, store):
c, _ = con
write_elements(
c,
"99490",
[_el("99490", 2026, "consent", "required", "YBM4IZUS", 1578)],
[],
)
write_elements(c, "99491", [_el("99491", 2026, "consent", "required")], [])
write_elements(c, "G0556", [_el("G0556", 2026, "telehealth", "yes")], [])
diffs, note = lineage._element_diffs(c, ["99490", "99491", "G0556"], store, 1)
assert note == ""
by_pair = {(d.type, d.value): d for d in diffs}
assert set(by_pair) == {("consent", "required"), ("telehealth", "yes")}
consent = by_pair[("consent", "required")]
assert consent.in_codes == ("99490", "99491")
assert consent.not_in_codes == ("G0556",)
assert consent.label == "CY2021 PFS final ¶1578" # from the anchored 99490 row
telehealth = by_pair[("telehealth", "yes")]
assert telehealth.in_codes == ("G0556",)
assert telehealth.not_in_codes == ("99490", "99491")
def test_no_diff_when_every_code_with_elements_shares_the_value(self, con):
c, _ = con
write_elements(c, "99490", [_el("99490", 2026, "consent", "required")], [])
write_elements(c, "99491", [_el("99491", 2026, "consent", "required")], [])
diffs, _ = lineage._element_diffs(c, ["99490", "99491"], None, 1)
assert diffs == ()
def test_only_the_newest_year_per_code_is_compared(self, con):
c, _ = con
write_elements(
c,
"99490",
[
_el("99490", 2020, "consent", "verbal"),
_el("99490", 2026, "consent", "required"),
],
[],
)
write_elements(c, "99491", [_el("99491", 2026, "consent", "verbal")], [])
diffs, _ = lineage._element_diffs(c, ["99490", "99491"], None, 1)
# 2020's "verbal" for 99490 must not be compared — only 2026's
# "required" — so the diff is required (99490) vs verbal (99491).
values = {d.value: d for d in diffs}
assert "verbal" not in {d.value for d in diffs if "99490" in d.in_codes}
assert values["required"].in_codes == ("99490",)
def test_capped_at_twelve_most_shared_first(self, con):
c, _ = con
codes = ["A0001", "A0002", "A0003"]
for code in codes:
rows = [_el(code, 2026, f"t{i}", f"v{i}") for i in range(20)]
# one shared value across all three codes, sorts first
rows.append(_el(code, 2026, "shared", "yes"))
write_elements(c, code, rows, [])
# a fourth code with none of the "t*" values, so every t* pair
# differs (present in the three, absent in the fourth)
write_elements(c, "A0004", [_el("A0004", 2026, "shared", "no")], [])
diffs, _ = lineage._element_diffs(c, codes + ["A0004"], None, 1)
assert len(diffs) == 12
assert diffs[0].in_codes == ("A0001", "A0002", "A0003")
class TestGuidance:
def test_cfr_first_deduped_and_capped(self, con, store):
c, _ = con
rows = [
_gd("CCM", "99490", "cfr", f"42 CFR 410.{i}", "YBM4IZUS", 686)
for i in range(6)
]
rows += [
_gd("CCM", "99490", "iom", f"100-04 ch.{i}", "YBM4IZUS", 686)
for i in range(6)
]
write_guidance(c, "CCM", rows)
out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
assert len(out) == 8
assert [g.kind for g in out[:6]] == ["cfr"] * 6
assert out[0].url == "https://www.ecfr.gov/current/title-42/section-410.0"
def test_dedupes_by_locator(self, con, store):
c, _ = con
write_guidance(
c,
"CCM",
[
_gd("CCM", "99490", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 1578),
_gd("CCM", "99491", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 686),
],
)
out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
assert len(out) == 1
def test_wide_families_are_skipped(self, con, store):
c, _ = con
write_guidance(
c, "APCM", [_gd("APCM", "G0556", "mln", "MLN 907166", "YBM4IZUS", 686)]
)
out = lineage._collect_guidance(c, store, ["APCM"], ["APCM"], 1)
assert out == ()
def test_iom_mln_have_no_url(self, con, store):
c, _ = con
write_guidance(
c,
"CCM",
[_gd("CCM", "99490", "iom", "100-04 ch.12 §30.6.4", "YBM4IZUS", 686)],
)
out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
assert out[0].url == ""
assert out[0].label == "CY2021 PFS final ¶686"
class TestPromptBlock:
def test_exact_text(self):
events = (
LineageEvent(
code="G2058",
year=2021,
kind="replaced_by",
from_codes=(),
to_codes=("99439",),
label="CY2021 PFS final ¶686",
item_key="YBM4IZUS",
p_id=686,
page=84547,
url="https://fr.test/2021-doc#p-686",
source="fr",
anchored=True,
note="",
),
LineageEvent(
code="99439",
year=2021,
kind="created",
from_codes=(),
to_codes=(),
label="CY2021 PFS final ¶686",
item_key="YBM4IZUS",
p_id=686,
page=84547,
url="https://fr.test/2021-doc#p-686",
source="fr",
anchored=True,
note="from CPT Editorial Panel",
),
)
diffs = (
ElementDiff(
type="consent",
value="required",
in_codes=("99490", "99491"),
not_in_codes=("G0556",),
label="CY2021 PFS final ¶1578",
item_key="YBM4IZUS",
p_id=1578,
),
)
guidance = (
GuidanceRef(
kind="cfr",
locator="42 CFR 410.78(a)(3)",
url="https://www.ecfr.gov/current/title-42/section-410.78",
label="CY2021 PFS final ¶1578",
item_key_src="YBM4IZUS",
p_id_src=1578,
),
)
ev = LineageEvidence(
codes=("99439", "99490", "99491", "G0556", "G2058"),
families=("CCM", "APCM"),
events=events,
element_diffs=diffs,
guidance=guidance,
)
assert ev.prompt_block() == (
"Lineage (dated events with anchors; cite the bracketed label "
"for any dated claim):\n"
"[CY2021 PFS final ¶686] 2021 replaced_by G2058 ( → 99439)\n"
"[CY2021 PFS final ¶686] 2021 created 99439 — "
"from CPT Editorial Panel\n"
"Element differences:\n"
"[CY2021 PFS final ¶1578] consent=required: "
"in 99490, 99491; not in G0556\n"
"Guidance:\n"
"[CY2021 PFS final ¶1578] 42 CFR 410.78(a)(3) — CFR"
)
def test_header_only_with_no_events(self):
ev = LineageEvidence(
codes=("99490",), families=(), events=(), element_diffs=(), guidance=()
)
assert ev.prompt_block() == (
"Lineage (dated events with anchors; cite the bracketed label "
"for any dated claim):"
)
def test_priority_kinds_always_kept_non_priority_trimmed(self):
priority = LineageEvent(
"99490", 2021, "created", (), (), "[L]", "K", 1, 0, "", "fr", True, ""
)
noisy = [
LineageEvent(
"99490",
2022 + i,
"revalued",
(),
(),
f"[L{i}]",
"K",
1,
0,
"",
"rvu",
True,
"",
)
for i in range(5)
]
ev = LineageEvidence(
codes=("99490",),
families=(),
events=(priority, *noisy),
element_diffs=(),
guidance=(),
max_prompt_rows=3,
)
rendered = ev.prompt_block()
assert "created" in rendered
assert rendered.count("revalued") == 2 # budget: 3 - 1 priority = 2 kept
def test_priority_kinds_exceed_cap_but_are_never_dropped(self):
priorities = [
LineageEvent(
"99490", 2020 + i, k, (), (), f"[L{i}]", "K", 1, 0, "", "fr", True, ""
)
for i, k in enumerate(["created", "replaces", "replaced_by", "deleted"])
]
ev = LineageEvidence(
codes=("99490",),
families=(),
events=tuple(priorities),
element_diffs=(),
guidance=(),
max_prompt_rows=2,
)
rendered = ev.prompt_block()
for kind in ("created", "replaces", "replaced_by", "deleted"):
assert kind in rendered
class TestPayload:
def test_json_serializable_and_shaped(self):
events = (
LineageEvent(
"99439",
2021,
"created",
(),
(),
"[L]",
"YBM4IZUS",
686,
84547,
"u",
"fr",
True,
"",
),
)
ev = LineageEvidence(
codes=("99439",),
families=("CCM",),
events=events,
element_diffs=(),
guidance=(),
elements_note="elements extracted for 1 of 3 codes",
)
payload = ev.payload()
raw = json.dumps(payload)
back = json.loads(raw)
assert back["type"] == "lineage"
assert back["codes"] == ["99439"]
assert back["families"] == ["CCM"]
assert back["events"][0]["kind"] == "created"
assert back["elements_note"] == "elements extracted for 1 of 3 codes"
def test_no_elements_note_key_when_empty(self):
ev = LineageEvidence(("99439",), (), (), (), ())
assert "elements_note" not in ev.payload()
class TestLineageEvidence:
def test_none_when_no_codes_detected(self):
assert (
lineage_evidence("what did commenters say about telehealth?", CFG) is None
)
def test_collapsed_events_and_labels_for_the_ccm_fixture(self, con, store):
c, path = con
write_events(
c,
"99490",
[
_ev(
"99490",
2015,
"adopted_cpt",
item_key="DE2VH9PD",
p_id=1251,
page=67716,
)
],
)
write_events(
c,
"G2058",
[
_ev(
"G2058",
2021,
"replaced_by",
to="99439",
item_key="YBM4IZUS",
p_id=686,
page=84547,
),
_ev("G2058", 2022, "disappeared", source="rvu", item_key="", p_id=0),
],
)
write_events(
c,
"99439",
[
_ev(
"99439", 2021, "created", item_key="YBM4IZUS", p_id=686, page=84547
),
_ev(
"99439",
2021,
"replaces",
frm="G2058",
item_key="YBM4IZUS",
p_id=1578,
page=84639,
),
# same group as the fr "created" row above — collapse must
# drop this rvu duplicate.
_ev(
"99439",
2021,
"created",
source="rvu",
item_key="",
p_id=0,
note="dupe",
),
],
)
c.close() # release the write handle before lineage_evidence opens read-only
cfg = replace(CFG, duckdb_replica=str(path))
det = Detection(
codes=("99439", "99490", "G2058"), families=(), explicit=(), wide=()
)
with patch("llm.lineage._store", return_value=store):
with patch("llm.lineage.detect_codes", return_value=det):
ev = lineage_evidence("irrelevant text", cfg)
assert ev is not None
assert [(e.year, e.code, e.kind, e.source) for e in ev.events] == [
(2015, "99490", "adopted_cpt", "fr"),
(2021, "99439", "created", "fr"),
(2021, "99439", "replaces", "fr"),
(2021, "G2058", "replaced_by", "fr"),
(2022, "G2058", "disappeared", "rvu"),
]
adopted, created, replaces, replaced_by, disappeared = ev.events
assert adopted.label == "CY2015 PFS final ¶1251"
assert created.label == "CY2021 PFS final ¶686"
assert replaces.label == "CY2021 PFS final ¶1578"
assert replaces.from_codes == ("G2058",)
assert replaced_by.to_codes == ("99439",)
assert disappeared.label == "PFS CY2022 RVU file"
assert created.url == "https://fr.test/2021-doc#p-686"
def test_on_demand_fallback_only_for_codes_without_rows_and_capped(
self, con, store
):
c, path = con
rvu_cols = (
"hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
"non_fac_total DOUBLE, year INTEGER"
)
c.execute(f"CREATE TABLE pfs.rvu ({rvu_cols})")
c.executemany(
"INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?)",
[
*[("00000", None, "filler", "A", 1.0, y) for y in range(2015, 2027)],
("99441", None, "d", "A", 1.0, 2020),
("99441", None, "d", "A", 1.0, 2021),
("99442", None, "d", "A", 1.0, 2020),
("99442", None, "d", "A", 1.0, 2021),
],
)
# 99490 has precomputed rows — must not trigger the fallback.
write_events(
c,
"99490",
[_ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251)],
)
c.close() # release the write handle before lineage_evidence opens read-only
cfg = replace(CFG, duckdb_replica=str(path), lineage_on_demand_max=1)
det = Detection(
codes=("99441", "99442", "99490"), families=(), explicit=(), wide=()
)
with patch("llm.lineage._store", return_value=store):
with patch("llm.lineage.detect_codes", return_value=det):
with patch(
"llm.lineage._on_demand_events", wraps=lineage._on_demand_events
) as spy:
ev = lineage_evidence("irrelevant text", cfg)
assert ev is not None
# budget is 1: only the first no-rows code (99441, sorted first)
# gets the fallback; 99442 and the precomputed 99490 do not.
assert spy.call_count == 1
assert spy.call_args.args[-1] == "99441"
codes_with_events = {e.code for e in ev.events}
assert "99441" in codes_with_events
assert "99442" not in codes_with_events
def test_never_raises_on_a_broken_replica(self, tmp_path):
bad = tmp_path / "not-a-duckdb-file"
bad.write_text("not a database")
cfg = replace(CFG, duckdb_replica=str(bad))
det = Detection(codes=("99490",), families=(), explicit=(), wide=())
with patch("llm.lineage.detect_codes", return_value=det):
assert lineage_evidence("irrelevant", cfg) is None
def test_elements_note_and_guidance_flow_through(self, con, store):
c, path = con
write_elements(c, "99490", [_el("99490", 2026, "consent", "required")], [])
write_guidance(
c,
"CCM",
[_gd("CCM", "99490", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 1578)],
)
c.close() # release the write handle before lineage_evidence opens read-only
cfg = replace(CFG, duckdb_replica=str(path))
det = Detection(codes=("99490",), families=("CCM",), explicit=(), wide=())
with patch("llm.lineage._store", return_value=store):
with patch("llm.lineage.detect_codes", return_value=det):
ev = lineage_evidence("irrelevant", cfg)
assert ev is not None
assert ev.elements_note == "elements extracted for 1 of 1 codes"
assert ev.element_diffs == ()
assert len(ev.guidance) == 1
assert ev.guidance[0].locator == "42 CFR 410.78(a)(3)"
class TestCptLabel:
def test_cpt_label_is_always_cpt_changes_year_note_holds_the_detail(self):
"""CPT-sourced events (cpt_changed, cpt_deleted) always label as
"CPT Changes {year}" — the note (e.g. "absent from CPT 2024")
is a separate field, not folded into the label."""
row = _ev(
"99213",
2024,
"cpt_deleted",
item_key="CPTED2024",
p_id=0,
source="cpt",
note="absent from CPT 2024",
)
event = lineage._to_lineage_event(None, row, 1)
assert event.label == "CPT Changes 2024"
assert event.note == "absent from CPT 2024"
line = lineage._event_line(event)
assert (
line == "[CPT Changes 2024] 2024 cpt_deleted 99213 — absent from CPT 2024"
)
assert (
line.count("absent from CPT 2024") == 1
) # not duplicated into the label too
def test_cpt_changed_label_ignores_any_note(self):
row = _ev(
"99213",
2022,
"cpt_changed",
item_key="CPTED2022",
p_id=0,
source="cpt",
note="some other detail",
)
event = lineage._to_lineage_event(None, row, 1)
assert event.label == "CPT Changes 2022"
class TestSourceLabelAndUrlDegrade:
def test_source_label_degrades_to_item_key_when_store_get_fails(self):
store = SimpleNamespace(get=lambda k: (_ for _ in ()).throw(KeyError(k)))
assert lineage._source_label(store, "UNKNOWN1", 5, 1) == "UNKNOWN1"
def test_source_label_degrades_to_question_mark_when_item_key_empty(self):
store = SimpleNamespace(get=lambda k: (_ for _ in ()).throw(KeyError(k)))
assert lineage._source_label(store, "", 5, 1) == "?"
def test_fr_url_degrades_to_empty_string_on_resolve_failure(self, store):
with patch("bib.frlink.resolve", side_effect=RuntimeError("boom")):
assert lineage._fr_url(store, "YBM4IZUS", 686, 1) == ""
def test_fr_url_empty_when_item_key_or_p_id_missing(self):
assert lineage._fr_url(None, "", 5, 1) == ""
assert lineage._fr_url(None, "K", 0, 1) == ""
class TestCachingAndScopedUrlResolution:
"""Ruling: Store.get results are cached per item_key (not just per
(item_key, p_id) — the same rule backs many events/diffs/guidance
rows), and FR urls are only resolved for the rows that actually
reach the prompt block."""
def test_urls_resolved_only_for_prompt_rows_labels_cached_per_item(
self, con, store
):
c, path = con
n_items = 10
rows = [
_ev(
"99999",
2000 + i,
"crosswalk", # non-priority, not an _EARLIEST_ONLY kind
item_key=f"K{i % n_items:02d}",
p_id=1000 + i,
source="fr",
anchored=True,
)
for i in range(60)
]
write_events(c, "99999", rows)
store.con.executemany(
"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)
],
)
c.close() # release the write handle before lineage_evidence opens read-only
cfg = replace(CFG, duckdb_replica=str(path), lineage_max_rows=25)
det = Detection(codes=("99999",), families=(), explicit=(), wide=())
real_get = store.get
seen_keys: list[str] = []
def counting_get(key):
seen_keys.append(key)
return real_get(key)
with patch("llm.lineage._store", return_value=store):
with patch("llm.lineage.detect_codes", return_value=det):
with patch.object(store, "get", side_effect=counting_get) as get_spy:
with patch("bib.frlink.resolve") as resolve_mock:
resolve_mock.return_value = SimpleNamespace(url="http://x")
ev = lineage_evidence("irrelevant", cfg)
assert ev is not None
assert len(ev.events) == 60 # payload carries every collapsed row
# exactly the 25 prompt-selected rows (budget: no priority kinds
# here, so the first 25 by year) get an FR url resolved.
assert resolve_mock.call_count == 25
with_url = [e for e in ev.events if e.url]
assert len(with_url) == 25
# every referenced item_key is fetched at most once, total calls
# bounded by the number of distinct item_keys (10), not events (60).
assert get_spy.call_count <= n_items
assert len(seen_keys) == len(set(seen_keys)) # never re-fetched
def _lineage_evidence(
events, max_prompt_rows=25, explicit=(), element_diffs=()
) -> LineageEvidence:
return LineageEvidence(
codes=(),
families=(),
events=tuple(events),
element_diffs=tuple(element_diffs),
guidance=(),
max_prompt_rows=max_prompt_rows,
explicit=tuple(explicit),
)
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) == []
def _insert_extra_anchor(store, key, year, p_id, text="extra text"):
store.con.execute(
"INSERT INTO items VALUES (?,?,?,?)",
(
key,
f"Medicare Program; CY {year} PFS Final Rule",
f"{year - 1}-11-01",
f"https://www.federalregister.gov/d/{year}-doc-{key}",
),
)
store.con.execute(
"INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
(key, p_id, 1000 + p_id, 1, text),
)
store.con.commit()
def _diff(item_key, p_id, label, type_="element", value="v") -> ElementDiff:
return ElementDiff(
type=type_,
value=value,
in_codes=("A",),
not_in_codes=("B",),
label=label,
item_key=item_key,
p_id=p_id,
)
class TestLineageSourcesRulingB10:
"""Ruling B10 — explicit-code-first ordering and element-diff
anchors appended to ``lineage_sources``' output."""
def test_explicit_code_sorts_first_even_when_oldest_and_non_priority(self, store):
# A (explicit, literally named in the question) has the oldest,
# non-priority-kind event; B and C (family-mates, not named) have
# newer, priority-kind events. Explicit-code membership must
# outrank both the priority-kind tier and the year — A still
# sorts first.
a = lineage._to_lineage_event(
store,
_ev("A", 2015, "revalued", item_key="DE2VH9PD", p_id=1251),
mtime=0,
)
b = lineage._to_lineage_event(
store,
_ev("B", 2021, "created", item_key="YBM4IZUS", p_id=686),
mtime=0,
)
c = lineage._to_lineage_event(
store,
_ev("C", 2021, "created", item_key="YBM4IZUS", p_id=1578),
mtime=0,
)
ev = _lineage_evidence([a, b, c], explicit=("A",))
out = lineage_sources(store, ev, max_items=8)
assert [s["label"] for s in out] == [a.label, b.label, c.label]
def test_explicit_code_survives_a_cap_that_would_otherwise_drop_it(self, store):
a = lineage._to_lineage_event(
store,
_ev("A", 2015, "revalued", item_key="DE2VH9PD", p_id=1251),
mtime=0,
)
b = lineage._to_lineage_event(
store,
_ev("B", 2021, "created", item_key="YBM4IZUS", p_id=686),
mtime=0,
)
ev = _lineage_evidence([a, b], explicit=("A",))
(out,) = lineage_sources(store, ev, max_items=1)
assert out["label"] == a.label
def test_element_diff_anchors_appended_after_events(self, store):
e = lineage._to_lineage_event(
store,
_ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
mtime=0,
)
diff = _diff("YBM4IZUS", 686, "CY2021 PFS final ¶686")
ev = _lineage_evidence([e], element_diffs=[diff])
out = lineage_sources(store, ev, max_items=8)
assert [s["label"] for s in out] == [e.label, diff.label]
assert out[1]["kind"] == "rule"
assert out[1]["item_key"] == "YBM4IZUS"
assert out[1]["p_id"] == "686"
def test_element_diff_anchors_capped_at_two_most_shared_first(self, store):
_insert_extra_anchor(store, "EX1", 2016, 10)
_insert_extra_anchor(store, "EX2", 2017, 11)
diffs = [
_diff("DE2VH9PD", 1251, "CY2015 PFS final ¶1251"), # most shared
_diff("EX1", 10, "CY2016 PFS final ¶10"),
_diff("EX2", 11, "CY2017 PFS final ¶11"), # third — dropped by the cap
]
ev = _lineage_evidence([], element_diffs=diffs)
out = lineage_sources(store, ev, max_items=8)
assert [s["label"] for s in out] == [diffs[0].label, diffs[1].label]
def test_element_diffs_still_added_past_max_items_up_to_the_hard_cap(self, store):
# 8 distinct event paragraphs exactly fill max_items=8; two more
# distinct element-diff paragraphs are still appended (the "add
# up to 2 even when max_items has no room" rule), reaching the
# absolute hard cap of 10, not 11.
events = []
for i in range(8):
key = f"EV{i}"
_insert_extra_anchor(store, key, 2000 + i, 100 + i)
events.append(
lineage._to_lineage_event(
store,
_ev("Z", 2000 + i, "created", item_key=key, p_id=100 + i),
mtime=0,
)
)
_insert_extra_anchor(store, "DX1", 2050, 500)
_insert_extra_anchor(store, "DX2", 2051, 501)
_insert_extra_anchor(store, "DX3", 2052, 502)
diffs = [
_diff("DX1", 500, "d1"),
_diff("DX2", 501, "d2"),
_diff("DX3", 502, "d3"), # third diff — dropped by the element cap of 2
]
ev = _lineage_evidence(events, explicit=("Z",), element_diffs=diffs)
out = lineage_sources(store, ev, max_items=8)
assert len(out) == 10
assert [s["label"] for s in out[8:]] == ["d1", "d2"]
def test_element_diff_with_no_anchor_is_skipped(self, store):
e = lineage._to_lineage_event(
store,
_ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
mtime=0,
)
rvu_only_diff = ElementDiff(
type="t",
value="v",
in_codes=("A",),
not_in_codes=("B",),
label="l",
item_key="",
p_id=0,
)
ev = _lineage_evidence([e], element_diffs=[rvu_only_diff])
out = lineage_sources(store, ev, max_items=8)
assert [s["label"] for s in out] == [e.label]
def test_element_diff_deduped_against_an_event_paragraph(self, store):
e = lineage._to_lineage_event(
store,
_ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
mtime=0,
)
same_paragraph_diff = _diff("DE2VH9PD", 1251, "some other label")
ev = _lineage_evidence([e], element_diffs=[same_paragraph_diff])
out = lineage_sources(store, ev, max_items=8)
assert len(out) == 1
assert out[0]["label"] == e.label
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):
from conf import ROOT
replica = ROOT / "data" / "replica" / "aco.ro.duckdb"
if not replica.exists():
pytest.skip(f"no live replica at {replica}")
cfg = replace(CFG, duckdb_replica=str(replica))
lineage_evidence("history of CCM coding and payment", cfg) # warm
times_ms = []
for _ in range(5):
start = time.perf_counter()
lineage_evidence("history of CCM coding and payment", cfg)
times_ms.append((time.perf_counter() - start) * 1000)
times_ms.sort()
median = times_ms[len(times_ms) // 2]
print(f"lineage_evidence median: {median:.2f} ms (all: {times_ms})")
assert median < 300, f"{median:.2f} ms over {times_ms}"