1036 lines
37 KiB
Python
1036 lines
37 KiB
Python
"""llm.evidence — valuation prompt block, SSE payload, code-cited rule sources."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import replace
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import duckdb
|
|
import pytest
|
|
|
|
from llm import evidence
|
|
from llm.config import LlmConfig
|
|
from llm.evidence import (
|
|
ValuationEvidence,
|
|
code_cited_sources,
|
|
manual_sources,
|
|
merge_sources,
|
|
valuation_evidence,
|
|
warm,
|
|
)
|
|
from pfs.codetables import ensure_tables
|
|
from pfs.valuation import ValuationRow
|
|
|
|
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",
|
|
valuation_years=4,
|
|
code_cited_per_code=2,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clear_replica_cache():
|
|
"""The replica handle is process-global; no test may inherit another's."""
|
|
evidence._REPLICA = None
|
|
yield
|
|
evidence._REPLICA = None
|
|
|
|
|
|
def _row(**kw) -> ValuationRow:
|
|
base = dict(
|
|
code="G0556",
|
|
description="Adv prim care mgmt lvl 1",
|
|
vintage="CY2026 final",
|
|
year=2026,
|
|
proposed=False,
|
|
status="A",
|
|
work=0.25,
|
|
pe_nf=0.22,
|
|
pe_f=0.06,
|
|
mp=0.02,
|
|
total_nf=0.49,
|
|
total_f=0.33,
|
|
cf=33.4009,
|
|
pay_nf=16.37,
|
|
pay_f=11.02,
|
|
label="[PFS CY2026 Addendum B]",
|
|
citation="90 FR 49266",
|
|
url="https://www.federalregister.gov/citation/90-FR-49266",
|
|
)
|
|
base.update(kw)
|
|
return ValuationRow(**base)
|
|
|
|
|
|
class TestEvidenceText:
|
|
def test_prompt_block_exact(self):
|
|
ev = ValuationEvidence(
|
|
codes=("G0556",), families=("APCM",), rows=(_row(),), unpriced=("Z9999",)
|
|
)
|
|
assert ev.prompt_block() == (
|
|
"Valuation (authoritative for RVUs, conversion factors and payments; "
|
|
"national unadjusted, no GPCI):\n"
|
|
'[PFS CY2026 Addendum B] G0556 "Adv prim care mgmt lvl 1" — status A; work 0.25; '
|
|
"PE 0.22 non-fac / 0.06 fac; MP 0.02; total 0.49 non-fac / 0.33 fac; "
|
|
"CF $33.4009 → payment $16.37 non-fac / $11.02 fac\n"
|
|
"Not priced in the indexed fee schedule: Z9999"
|
|
)
|
|
|
|
def test_prompt_block_labels_the_cf_and_an_unpaid_status(self):
|
|
ev = ValuationEvidence(
|
|
("G0559",),
|
|
(),
|
|
(
|
|
_row(
|
|
code="G0559",
|
|
status="B",
|
|
status_note="bundled — no separate PFS payment",
|
|
cf_note="non-APM standard CF; QP CF $33.5675",
|
|
pay_nf=None,
|
|
pay_f=None,
|
|
),
|
|
),
|
|
(),
|
|
)
|
|
line = ev.prompt_block().splitlines()[1]
|
|
assert "— status B; bundled — no separate PFS payment; work" in line
|
|
assert "CF $33.4009 [non-APM standard CF; QP CF $33.5675] → payment" in line
|
|
assert "payment n/a non-fac / n/a fac" in line
|
|
|
|
def test_prompt_block_null_components(self):
|
|
ev = ValuationEvidence(
|
|
("G9999",),
|
|
(),
|
|
(
|
|
_row(
|
|
code="G9999",
|
|
status="I",
|
|
work=None,
|
|
pe_nf=None,
|
|
pe_f=None,
|
|
mp=None,
|
|
total_nf=None,
|
|
total_f=None,
|
|
pay_nf=None,
|
|
pay_f=None,
|
|
),
|
|
),
|
|
(),
|
|
)
|
|
line = ev.prompt_block().splitlines()[1]
|
|
assert (
|
|
"status I; work n/a; PE n/a non-fac / n/a fac" in line
|
|
and "payment n/a" in line
|
|
)
|
|
|
|
def test_payload_shape(self):
|
|
ev = ValuationEvidence(
|
|
("G0556",),
|
|
("APCM",),
|
|
(
|
|
_row(),
|
|
_row(
|
|
year=2027,
|
|
proposed=True,
|
|
vintage="CY2027 proposed",
|
|
label="[CY2027 NPRM Addendum B]",
|
|
citation="91 FR 43842",
|
|
url="https://www.federalregister.gov/citation/91-FR-43842",
|
|
),
|
|
),
|
|
(),
|
|
)
|
|
p = ev.payload()
|
|
assert (
|
|
p["type"] == "valuation"
|
|
and p["codes"] == ["G0556"]
|
|
and p["families"] == ["APCM"]
|
|
)
|
|
assert p["rows"][0]["pay_nf"] == 16.37 and p["rows"][1]["proposed"] is True
|
|
assert p["rows"][0]["status_note"] == "" and p["rows"][0]["cf_note"] == ""
|
|
assert p["provenance"] == [
|
|
{
|
|
"label": "[PFS CY2026 Addendum B]",
|
|
"vintage": "CY2026 final",
|
|
"citation": "90 FR 49266",
|
|
"url": "https://www.federalregister.gov/citation/90-FR-49266",
|
|
},
|
|
{
|
|
"label": "[CY2027 NPRM Addendum B]",
|
|
"vintage": "CY2027 proposed",
|
|
"citation": "91 FR 43842",
|
|
"url": "https://www.federalregister.gov/citation/91-FR-43842",
|
|
},
|
|
]
|
|
assert p["unpriced"] == []
|
|
|
|
|
|
class TestValuationEvidence:
|
|
def test_none_when_no_codes(self):
|
|
assert (
|
|
valuation_evidence("what did commenters say about telehealth?", CFG) is None
|
|
)
|
|
|
|
@patch("llm.evidence.valuation", return_value=([_row()], []))
|
|
@patch("llm.evidence.duckdb.connect")
|
|
def test_detects_and_queries_replica_read_only(self, mock_connect, mock_val):
|
|
ev = valuation_evidence("How is APCM valued?", CFG)
|
|
mock_connect.assert_called_once_with(
|
|
"/nonexistent/aco.ro.duckdb", read_only=True
|
|
)
|
|
assert mock_val.call_args.args[1] == ["G0556", "G0557", "G0558"]
|
|
assert mock_val.call_args.kwargs == {"years": 4}
|
|
assert (
|
|
ev is not None
|
|
and ev.codes == ("G0556", "G0557", "G0558")
|
|
and ev.families == ("APCM",)
|
|
)
|
|
mock_connect.return_value.close.assert_not_called() # cached, not closed
|
|
|
|
@patch("llm.evidence.duckdb.connect", side_effect=OSError("no replica"))
|
|
def test_unopenable_replica_yields_none(self, _c, caplog):
|
|
assert valuation_evidence("G0556?", CFG) is None
|
|
assert "valuation evidence skipped" in caplog.text
|
|
|
|
@patch("llm.evidence.valuation", return_value=([_row()], []))
|
|
@patch("llm.evidence.duckdb.connect")
|
|
def test_replica_connection_is_reused(self, mock_connect, _val):
|
|
valuation_evidence("G0556?", CFG)
|
|
valuation_evidence("G0557?", CFG)
|
|
mock_connect.assert_called_once()
|
|
mock_connect.return_value.close.assert_not_called()
|
|
|
|
@patch("llm.evidence.valuation", return_value=([_row()], []))
|
|
@patch("llm.evidence.duckdb.connect")
|
|
def test_republished_replica_reopens_and_closes_the_old_handle(
|
|
self, mock_connect, _val, monkeypatch
|
|
):
|
|
first, second = MagicMock(), MagicMock()
|
|
mock_connect.side_effect = [first, second]
|
|
mtimes = iter([1, 2])
|
|
monkeypatch.setattr(evidence, "_mtime", lambda _p: next(mtimes))
|
|
valuation_evidence("G0556?", CFG)
|
|
valuation_evidence("G0557?", CFG)
|
|
assert mock_connect.call_count == 2
|
|
first.close.assert_called_once()
|
|
second.close.assert_not_called()
|
|
|
|
@patch("llm.evidence.valuation", return_value=([], []))
|
|
@patch("llm.evidence.duckdb.connect")
|
|
def test_many_codes_shorten_the_window(self, _c, mock_val):
|
|
valuation_evidence("How is CCM valued?", CFG) # six-code family
|
|
assert mock_val.call_args.kwargs == {"years": 2}
|
|
|
|
|
|
class TestWarm:
|
|
"""#699 ruling B3: warm the replica connection ahead of the first
|
|
real request — at API startup and again (a no-op once cached) at the
|
|
top of every ``valuation_evidence`` call."""
|
|
|
|
def test_warm_opens_once_and_is_a_noop_on_replay(self, tmp_path):
|
|
path = tmp_path / "aco.ro.duckdb"
|
|
path.touch() # warm() only needs the file to exist, not be a real db
|
|
cfg = replace(CFG, duckdb_replica=str(path))
|
|
with patch("llm.evidence.duckdb.connect") as mock_connect:
|
|
mock_connect.return_value = MagicMock()
|
|
warm(cfg)
|
|
warm(cfg)
|
|
mock_connect.assert_called_once_with(str(path), read_only=True)
|
|
|
|
def test_warm_is_a_noop_when_the_replica_file_is_missing(self):
|
|
with patch("llm.evidence.duckdb.connect") as mock_connect:
|
|
warm(CFG) # CFG.duckdb_replica ("/nonexistent/...") doesn't exist
|
|
mock_connect.assert_not_called()
|
|
|
|
def test_warm_swallows_connect_failures(self, tmp_path, caplog):
|
|
path = tmp_path / "aco.ro.duckdb"
|
|
path.touch()
|
|
cfg = replace(CFG, duckdb_replica=str(path))
|
|
with patch("llm.evidence.duckdb.connect", side_effect=OSError("boom")):
|
|
warm(cfg) # must not raise
|
|
assert "replica warm-up skipped" in caplog.text
|
|
|
|
@patch("llm.evidence.detect_codes")
|
|
@patch("llm.evidence.warm")
|
|
def test_valuation_evidence_warms_before_detecting_codes(
|
|
self, mock_warm, mock_detect
|
|
):
|
|
order: list[str] = []
|
|
mock_warm.side_effect = lambda cfg: order.append("warm")
|
|
|
|
def _detect(question):
|
|
order.append("detect")
|
|
return SimpleNamespace(codes=())
|
|
|
|
mock_detect.side_effect = _detect
|
|
assert valuation_evidence("anything", CFG) is None
|
|
assert order == ["warm", "detect"]
|
|
mock_warm.assert_called_once_with(CFG)
|
|
|
|
|
|
class TestConnectRefreshesFamilies:
|
|
"""#699: derived families (``pfs.code_family``) live only on the
|
|
replica — ``_connect`` is the chat's one hook to pick them up, on
|
|
every open and every reopen (a republished replica)."""
|
|
|
|
@patch("llm.evidence.refresh_from")
|
|
@patch("llm.evidence.duckdb.connect")
|
|
def test_connect_refreshes_families_on_open_and_reopen(
|
|
self, mock_connect, mock_refresh, monkeypatch
|
|
):
|
|
mtimes = iter([1, 1, 2])
|
|
monkeypatch.setattr(evidence, "_mtime", lambda _p: next(mtimes))
|
|
mock_connect.side_effect = [MagicMock(), MagicMock()]
|
|
|
|
con1 = evidence._connect("/x")
|
|
assert mock_refresh.call_count == 1
|
|
mock_refresh.assert_called_with(con1)
|
|
|
|
con2 = evidence._connect("/x") # same (path, mtime) — cached, no refresh
|
|
assert con2 is con1
|
|
assert mock_refresh.call_count == 1
|
|
|
|
con3 = evidence._connect("/x") # mtime changed — reopen, refresh again
|
|
assert mock_refresh.call_count == 2
|
|
mock_refresh.assert_called_with(con3)
|
|
|
|
@patch("llm.evidence.refresh_from", side_effect=RuntimeError("boom"))
|
|
@patch("llm.evidence.duckdb.connect")
|
|
def test_refresh_failure_does_not_break_the_handle(
|
|
self, _mock_connect, _mock_refresh, caplog
|
|
):
|
|
con = evidence._connect("/x")
|
|
assert con is not None
|
|
assert "family refresh skipped" in caplog.text
|
|
|
|
|
|
RVU_COLS = (
|
|
"hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
|
|
"work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, "
|
|
"non_fac_total DOUBLE, fac_total DOUBLE, conv_factor DOUBLE, year INTEGER"
|
|
)
|
|
PROPOSED_COLS = (
|
|
"hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
|
|
"work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, "
|
|
"cms_rule_id VARCHAR"
|
|
)
|
|
|
|
|
|
class TestConcurrentChats:
|
|
"""/chat streams run in Starlette's threadpool, so several turns share
|
|
the cached replica handle. Without a per-call cursor their result sets
|
|
cross-talk and the broad except swallows it as "no valuation"."""
|
|
|
|
@pytest.fixture
|
|
def replica(self, tmp_path):
|
|
db = tmp_path / "aco.ro.duckdb"
|
|
con = duckdb.connect(str(db))
|
|
con.execute("CREATE SCHEMA pfs")
|
|
con.execute(f"CREATE TABLE pfs.rvu ({RVU_COLS})")
|
|
con.execute(f"CREATE TABLE pfs.rvu_proposed ({PROPOSED_COLS})")
|
|
# two unfamilied codes with different row counts, so a swapped
|
|
# result set is visible as a wrong count
|
|
con.executemany(
|
|
"INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
[
|
|
(
|
|
"99213",
|
|
None,
|
|
"Office visit",
|
|
"A",
|
|
1.0,
|
|
1.0,
|
|
0.5,
|
|
0.05,
|
|
2.05,
|
|
1.55,
|
|
30.0,
|
|
2026,
|
|
),
|
|
*[
|
|
(
|
|
"99214",
|
|
None,
|
|
"Office visit",
|
|
"A",
|
|
1.5,
|
|
1.0,
|
|
0.5,
|
|
0.05,
|
|
2.55,
|
|
2.05,
|
|
30.0,
|
|
y,
|
|
)
|
|
for y in (2024, 2025, 2026)
|
|
],
|
|
],
|
|
)
|
|
con.close()
|
|
return db
|
|
|
|
def test_parallel_turns_share_one_handle_without_cross_talk(self, replica):
|
|
cfg = replace(CFG, duckdb_replica=str(replica))
|
|
real_connect, opens = duckdb.connect, []
|
|
|
|
def counting_connect(*a, **kw):
|
|
opens.append(a)
|
|
return real_connect(*a, **kw)
|
|
|
|
expected = {"99213": 1, "99214": 3}
|
|
questions = ["What does 99213 pay?", "What does 99214 pay?"]
|
|
|
|
def one(i):
|
|
q = questions[i % 2]
|
|
ev = valuation_evidence(q, cfg)
|
|
assert ev is not None, "valuation dropped under concurrency"
|
|
(code,) = ev.codes
|
|
return code, len(ev.rows)
|
|
|
|
with patch("llm.evidence.duckdb.connect", counting_connect):
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
results = list(pool.map(one, range(160)))
|
|
|
|
assert all(n == expected[code] for code, n in results), sorted(set(results))
|
|
assert len(opens) == 1 # one handle, shared
|
|
|
|
def test_cursor_is_closed_but_the_handle_is_kept(self, replica):
|
|
cfg = replace(CFG, duckdb_replica=str(replica))
|
|
assert valuation_evidence("99213?", cfg) is not None
|
|
assert evidence._REPLICA is not None
|
|
# the cached parent still answers — it was never closed
|
|
assert valuation_evidence("99214?", cfg) is not None
|
|
|
|
|
|
def _engine(rows):
|
|
engine = MagicMock()
|
|
conn = engine.begin.return_value.__enter__.return_value
|
|
conn.execute.return_value.fetchall.return_value = rows
|
|
return engine, conn
|
|
|
|
|
|
def _rule_md(p_id, codes, date="2024-12-09"):
|
|
return {
|
|
"kind": "rule",
|
|
"item_key": "R1",
|
|
"p_id": str(p_id),
|
|
"page": "97710",
|
|
"ordinal": "1",
|
|
"html_url": "https://www.federalregister.gov/d/2024-25382",
|
|
"fr_volume": "89",
|
|
"date": date,
|
|
"title": "CY2025 PFS final rule",
|
|
"codes": codes,
|
|
}
|
|
|
|
|
|
def _comment_md(seq, codes, item_key="C1", date="2024-12-09"):
|
|
return {
|
|
"kind": "comment",
|
|
"item_key": item_key,
|
|
"seq": str(seq),
|
|
"comment_id": "CMS-2026-2377-3438",
|
|
"date": date,
|
|
"codes": codes,
|
|
"families": "",
|
|
}
|
|
|
|
|
|
def _corpus_md(seq, codes="", families="", item_key="CPT-ED", date="2024-12-09"):
|
|
return {
|
|
"kind": "corpus",
|
|
"item_key": item_key,
|
|
"seq": str(seq),
|
|
"url": "https://example.org/cpt.pdf",
|
|
"title": "CPT Manual",
|
|
"date": date,
|
|
"codes": codes,
|
|
"families": families,
|
|
}
|
|
|
|
|
|
def _multi_engine(by_code=None, by_family=None):
|
|
"""A fake engine whose rows depend on which collection *and* which of
|
|
the two SQL statements (codes vs. families) is being executed —
|
|
distinguished by which bind parameter the call carries."""
|
|
engine = MagicMock()
|
|
conn = engine.begin.return_value.__enter__.return_value
|
|
by_code = by_code or {}
|
|
by_family = by_family or {}
|
|
|
|
def _execute(_sql, params):
|
|
collection = params["collection"]
|
|
table = by_code if "codes" in params else by_family
|
|
result = MagicMock()
|
|
result.fetchall.return_value = table.get(collection, [])
|
|
return result
|
|
|
|
conn.execute.side_effect = _execute
|
|
return engine, conn
|
|
|
|
|
|
class TestCodeCitedSources:
|
|
def test_sql_and_per_code_cap_and_dedupe(self):
|
|
engine, conn = _engine(
|
|
[
|
|
("para one G0556", _rule_md(1, "G0556")),
|
|
("para two G0556", _rule_md(2, "G0556")),
|
|
("para three G0556", _rule_md(3, "G0556")), # over the cap for G0556
|
|
(
|
|
"para four G0557",
|
|
_rule_md(4, "G0557 G0556"),
|
|
), # G0557 still under cap
|
|
("para one again", _rule_md(1, "G0556")), # duplicate (item_key, p_id)
|
|
]
|
|
)
|
|
out = code_cited_sources(engine, ["G0556", "G0557"], per_code=2)
|
|
sql = str(conn.execute.call_args.args[0])
|
|
assert "string_to_array" in sql and "langchain_pg_collection" in sql
|
|
assert "PARTITION BY" in sql
|
|
# the collection is a scalar subquery, not a join to filter after
|
|
assert "JOIN langchain_pg_collection" not in sql
|
|
assert "e.collection_id = (SELECT uuid FROM langchain_pg_collection" in sql
|
|
assert conn.execute.call_args.args[1]["collection"] == "rules"
|
|
assert conn.execute.call_args.args[1]["codes"] == ["G0556", "G0557"]
|
|
assert conn.execute.call_args.args[1]["window"] == 4
|
|
assert [s["snippet"] for s in out] == [
|
|
"para one G0556",
|
|
"para two G0556",
|
|
"para four G0557",
|
|
]
|
|
assert all(s["kind"] == "rule" and s["score"] == 0.0 for s in out)
|
|
|
|
def test_snippets_are_capped(self):
|
|
engine, _ = _engine([("x" * 900, _rule_md(1, "G0556"))])
|
|
(out,) = code_cited_sources(engine, ["G0556"], per_code=2)
|
|
assert len(out["snippet"]) == 250
|
|
|
|
def test_engine_error_yields_empty(self, caplog):
|
|
engine = MagicMock()
|
|
engine.begin.side_effect = RuntimeError("pg down")
|
|
assert code_cited_sources(engine, ["G0556"], per_code=2) == []
|
|
assert "code-cited sources skipped" in caplog.text
|
|
|
|
def test_no_codes_returns_empty_without_querying(self):
|
|
engine = MagicMock()
|
|
assert code_cited_sources(engine, [], per_code=2) == []
|
|
engine.begin.assert_not_called()
|
|
|
|
|
|
class TestCodeCitedSourcesMultiCollection:
|
|
def test_concatenated_rules_then_comments_then_corpus(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={
|
|
"rules": [("rule chunk", _rule_md(1, "G0556"))],
|
|
"comments": [("comment chunk", _comment_md(1, "G0556"))],
|
|
"corpus": [("corpus chunk", _corpus_md(1, "G0556"))],
|
|
}
|
|
)
|
|
out = code_cited_sources(
|
|
engine, ["G0556"], per_code=2, collections=("rules", "comments", "corpus")
|
|
)
|
|
assert [s["kind"] for s in out] == ["rule", "comment", "corpus"]
|
|
|
|
def test_per_collection_cap_applies_independently(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={
|
|
"rules": [
|
|
("r1", _rule_md(1, "G0556")),
|
|
("r2", _rule_md(2, "G0556")),
|
|
("r3", _rule_md(3, "G0556")), # over the per-collection cap
|
|
],
|
|
"comments": [
|
|
("c1", _comment_md(1, "G0556")),
|
|
("c2", _comment_md(2, "G0556")),
|
|
],
|
|
}
|
|
)
|
|
out = code_cited_sources(
|
|
engine, ["G0556"], per_code=2, collections=("rules", "comments")
|
|
)
|
|
# Capped to 2/collection, then round-robin interleaved across
|
|
# collections (rules, comments), not simply concatenated.
|
|
assert [s["snippet"] for s in out] == ["r1", "c1", "r2", "c2"]
|
|
|
|
def test_dedupe_uses_seq_for_non_rule_kinds(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={
|
|
"comments": [
|
|
("c1", _comment_md(1, "G0556", item_key="C1")),
|
|
("c1 dup", _comment_md(1, "G0556", item_key="C1")), # same seq
|
|
("c2", _comment_md(2, "G0556", item_key="C1")), # different seq
|
|
]
|
|
}
|
|
)
|
|
out = code_cited_sources(
|
|
engine, ["G0556"], per_code=5, collections=("comments",)
|
|
)
|
|
assert [s["snippet"] for s in out] == ["c1", "c2"]
|
|
|
|
|
|
class TestFamilyCitedSources:
|
|
def test_skipped_entirely_when_no_families(self):
|
|
engine, conn = _multi_engine(by_code={"rules": [("r1", _rule_md(1, "G0556"))]})
|
|
code_cited_sources(engine, ["G0556"], per_code=2, collections=("rules",))
|
|
assert conn.execute.call_count == 1 # codes query only, no family round trip
|
|
|
|
def test_sql_shape_and_params(self):
|
|
engine, conn = _multi_engine(by_family={"rules": []})
|
|
code_cited_sources(
|
|
engine, [], per_code=2, collections=("rules",), families=["CCM"]
|
|
)
|
|
sql = str(conn.execute.call_args.args[0])
|
|
assert "&&" in sql
|
|
# one chunk per (family, item_key), then the survivors ranked per family
|
|
assert sql.count("PARTITION BY") == 2
|
|
assert "PARTITION BY w.family, e.cmetadata->>'item_key'" in sql
|
|
assert "PARTITION BY family" in sql
|
|
assert conn.execute.call_args.args[1] == {
|
|
"collection": "rules",
|
|
"families": ["CCM"],
|
|
"window": 4,
|
|
}
|
|
|
|
def test_called_once_per_collection_with_same_params(self):
|
|
engine, conn = _multi_engine(
|
|
by_family={"rules": [], "comments": [], "corpus": []}
|
|
)
|
|
code_cited_sources(
|
|
engine,
|
|
[],
|
|
per_code=3,
|
|
collections=("rules", "comments", "corpus"),
|
|
families=["CCM"],
|
|
)
|
|
assert conn.execute.call_count == 3
|
|
for call, collection in zip(
|
|
conn.execute.call_args_list, ("rules", "comments", "corpus")
|
|
):
|
|
assert call.args[1] == {
|
|
"collection": collection,
|
|
"families": ["CCM"],
|
|
"window": 6,
|
|
}
|
|
|
|
def test_family_rows_appended_after_code_rows(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={"rules": [("code hit", _rule_md(1, "G0556"))]},
|
|
by_family={
|
|
"corpus": [
|
|
(
|
|
"CCM guideline",
|
|
_corpus_md(1, codes="", families="CCM", item_key="CPT-ED"),
|
|
)
|
|
]
|
|
},
|
|
)
|
|
out = code_cited_sources(
|
|
engine,
|
|
["G0556"],
|
|
per_code=2,
|
|
collections=("rules", "corpus"),
|
|
families=["CCM"],
|
|
)
|
|
assert [s["snippet"] for s in out] == ["code hit", "CCM guideline"]
|
|
assert out[1]["kind"] == "corpus"
|
|
|
|
def test_family_dedupes_against_a_code_cited_chunk(self):
|
|
md = _rule_md(1, "G0556")
|
|
md["families"] = "APCM"
|
|
engine, _conn = _multi_engine(
|
|
by_code={"rules": [("shared", md)]}, by_family={"rules": [("shared", md)]}
|
|
)
|
|
out = code_cited_sources(
|
|
engine, ["G0556"], per_code=2, collections=("rules",), families=["APCM"]
|
|
)
|
|
assert [s["snippet"] for s in out] == ["shared"]
|
|
|
|
|
|
class TestMaxTotal:
|
|
"""``max_total`` interleaves code/family rows round-robin across
|
|
collections, then truncates (#budget)."""
|
|
|
|
def test_interleaves_three_collections_round_robin(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={
|
|
"rules": [("r1", _rule_md(1, "G0556")), ("r2", _rule_md(2, "G0556"))],
|
|
"comments": [
|
|
("c1", _comment_md(1, "G0556")),
|
|
("c2", _comment_md(2, "G0556")),
|
|
],
|
|
"corpus": [
|
|
("x1", _corpus_md(1, codes="G0556", item_key="X1")),
|
|
("x2", _corpus_md(2, codes="G0556", item_key="X2")),
|
|
],
|
|
}
|
|
)
|
|
out = code_cited_sources(
|
|
engine,
|
|
["G0556"],
|
|
per_code=5,
|
|
collections=("rules", "comments", "corpus"),
|
|
max_total=0,
|
|
)
|
|
assert [s["snippet"] for s in out] == ["r1", "c1", "x1", "r2", "c2", "x2"]
|
|
|
|
def test_family_rows_interleaved_then_appended_after_code_rows(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={
|
|
"rules": [("r1", _rule_md(1, "G0556"))],
|
|
"comments": [("c1", _comment_md(1, "G0556"))],
|
|
},
|
|
by_family={
|
|
"rules": [
|
|
(
|
|
"f-rules",
|
|
_corpus_md(9, codes="", families="CCM", item_key="FR"),
|
|
)
|
|
],
|
|
"comments": [
|
|
(
|
|
"f-comments",
|
|
_corpus_md(10, codes="", families="CCM", item_key="FC"),
|
|
)
|
|
],
|
|
},
|
|
)
|
|
out = code_cited_sources(
|
|
engine,
|
|
["G0556"],
|
|
per_code=5,
|
|
collections=("rules", "comments"),
|
|
families=["CCM"],
|
|
max_total=0,
|
|
)
|
|
assert [s["snippet"] for s in out] == ["r1", "c1", "f-rules", "f-comments"]
|
|
|
|
def test_truncates_to_max_total(self):
|
|
engine, _conn = _multi_engine(
|
|
by_code={
|
|
"rules": [("r1", _rule_md(1, "G0556")), ("r2", _rule_md(2, "G0556"))],
|
|
"comments": [
|
|
("c1", _comment_md(1, "G0556")),
|
|
("c2", _comment_md(2, "G0556")),
|
|
],
|
|
}
|
|
)
|
|
out = code_cited_sources(
|
|
engine,
|
|
["G0556"],
|
|
per_code=5,
|
|
collections=("rules", "comments"),
|
|
max_total=3,
|
|
)
|
|
assert [s["snippet"] for s in out] == ["r1", "c1", "r2"]
|
|
|
|
def test_max_total_zero_or_negative_is_unlimited(self):
|
|
rows = [(f"r{i}", _rule_md(i, "G0556")) for i in range(1, 15)]
|
|
engine, _conn = _multi_engine(by_code={"rules": rows})
|
|
for unlimited in (0, -1):
|
|
out = code_cited_sources(
|
|
engine,
|
|
["G0556"],
|
|
per_code=20,
|
|
collections=("rules",),
|
|
max_total=unlimited,
|
|
)
|
|
assert len(out) == 14
|
|
|
|
def test_default_max_total_is_twelve(self):
|
|
rows = [(f"r{i}", _rule_md(i, "G0556")) for i in range(1, 15)]
|
|
engine, _conn = _multi_engine(by_code={"rules": rows})
|
|
out = code_cited_sources(engine, ["G0556"], per_code=20, collections=("rules",))
|
|
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}]
|
|
b = [{"label": "Y", "score": 0.0}, {"label": "Z", "score": 0.0}]
|
|
assert [s["label"] for s in merge_sources(a, b)] == ["X", "Y", "Z"]
|
|
assert merge_sources(a, b)[1]["score"] == 0.8
|
|
|
|
|
|
def _cpt_section_row(con, year, item_key, sec_id, title, path_key, guideline):
|
|
con.execute(
|
|
"INSERT INTO pfs.cpt_section VALUES (?,?,?,?,?,?,?,?,?,?)",
|
|
[
|
|
year,
|
|
item_key,
|
|
sec_id,
|
|
3,
|
|
title,
|
|
path_key.split(" > "),
|
|
path_key,
|
|
"",
|
|
"",
|
|
guideline,
|
|
],
|
|
)
|
|
|
|
|
|
def _family_row(con, key, note, code="99490"):
|
|
con.execute(
|
|
"INSERT INTO pfs.code_family VALUES (?,?,?,?,?,?,?,?,?)",
|
|
[key, key, code, "member", None, None, "", 0, note],
|
|
)
|
|
|
|
|
|
CCM_PATH = "Evaluation and Management > Care Management Services > Chronic Care Management Services"
|
|
PARENT_PATH = "Surgery > General"
|
|
LEAF_PATH = "Surgery > General > Leaf With No Guideline"
|
|
|
|
|
|
class TestManualSources:
|
|
"""``manual_sources`` — the CPT manual's own guideline text for a
|
|
detected family's heading, straight from the DuckDB replica (#691
|
|
Ruling B4: never a pgvector query — the CPT chunks' ``section``
|
|
metadata is junk EPUB headings, not the family heading)."""
|
|
|
|
@pytest.fixture
|
|
def cpt_replica(self):
|
|
con = duckdb.connect(":memory:")
|
|
ensure_tables(con)
|
|
# CCM: two editions have the heading's guideline text; the newer
|
|
# (2024) one must win over the older (2022) one.
|
|
_family_row(con, "CCM", CCM_PATH)
|
|
_cpt_section_row(
|
|
con,
|
|
2022,
|
|
"OLDED001",
|
|
"S0",
|
|
"Chronic Care Management Services",
|
|
CCM_PATH,
|
|
"2022 text: non-face-to-face care management services.",
|
|
)
|
|
_cpt_section_row(
|
|
con,
|
|
2024,
|
|
"GQGTPGYV",
|
|
"S1",
|
|
"Chronic Care Management Services",
|
|
CCM_PATH,
|
|
"Chronic care management services are non-face-to-face "
|
|
"services provided to a patient with two or more chronic conditions.",
|
|
)
|
|
# PARENTONLY: the leaf heading has no guideline of its own, but
|
|
# its one-level parent does — the fallback must pick it up.
|
|
_family_row(con, "PARENTONLY", LEAF_PATH)
|
|
_cpt_section_row(
|
|
con, 2024, "LEAFED001", "S2", "Leaf With No Guideline", LEAF_PATH, ""
|
|
)
|
|
_cpt_section_row(
|
|
con,
|
|
2023,
|
|
"PARED0001",
|
|
"S3",
|
|
"General",
|
|
PARENT_PATH,
|
|
"General guidance text for the parent heading.",
|
|
)
|
|
# NONOTE: a real family row, but its note is empty.
|
|
_family_row(con, "NONOTE", "", code="99999")
|
|
yield con
|
|
con.close()
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _fake_store(self):
|
|
items = {
|
|
"GQGTPGYV": SimpleNamespace(
|
|
date_published="2023-11-16",
|
|
url="https://ebooks.ama-assn.org/cpt2024",
|
|
)
|
|
}
|
|
|
|
def _get(key):
|
|
if key in items:
|
|
return items[key]
|
|
raise KeyError(key)
|
|
|
|
store = MagicMock()
|
|
store.get.side_effect = _get
|
|
with patch("llm.evidence._store", return_value=store):
|
|
yield store
|
|
|
|
def test_newest_edition_with_guideline_wins(self, cpt_replica):
|
|
out = manual_sources(cpt_replica, ["CCM"])
|
|
assert len(out) == 1
|
|
s = out[0]
|
|
assert s["kind"] == "corpus"
|
|
assert s["title"] == "CPT 2024 — Chronic Care Management Services"
|
|
assert s["item_key"] == "GQGTPGYV" and s["seq"] == "S1"
|
|
assert s["section"] == "Chronic Care Management Services"
|
|
assert s["date"] == "2023-11-16"
|
|
assert s["url"] == "https://ebooks.ama-assn.org/cpt2024"
|
|
assert s["snippet"].startswith("Chronic care management services")
|
|
|
|
def test_parent_fallback_when_leaf_has_no_guideline(self, cpt_replica):
|
|
out = manual_sources(cpt_replica, ["PARENTONLY"])
|
|
assert len(out) == 1
|
|
s = out[0]
|
|
assert s["title"] == "CPT 2023 — General"
|
|
assert s["item_key"] == "PARED0001" and s["seq"] == "S3"
|
|
assert "General guidance text" in s["snippet"]
|
|
# unresolved bib item (not in the fake store) — falls back to
|
|
# the plain edition year and an empty url.
|
|
assert s["date"] == "2023-01-01" and s["url"] == ""
|
|
|
|
def test_family_with_empty_note_yields_nothing(self, cpt_replica):
|
|
assert manual_sources(cpt_replica, ["NONOTE"]) == []
|
|
|
|
def test_family_with_no_family_row_yields_nothing(self, cpt_replica):
|
|
assert manual_sources(cpt_replica, ["GHOST"]) == []
|
|
|
|
def test_no_families_yields_nothing_without_querying(self, cpt_replica):
|
|
assert manual_sources(cpt_replica, []) == []
|
|
|
|
def test_per_family_cap(self, cpt_replica):
|
|
complex_path = (
|
|
"Evaluation and Management > Care Management Services > "
|
|
"Complex Chronic Care Management Services"
|
|
)
|
|
_family_row(cpt_replica, "CCM", complex_path, code="99487")
|
|
_cpt_section_row(
|
|
cpt_replica,
|
|
2024,
|
|
"GQGTPGYV",
|
|
"S4",
|
|
"Complex Chronic Care Management Services",
|
|
complex_path,
|
|
"Complex chronic care management guideline text.",
|
|
)
|
|
assert len(manual_sources(cpt_replica, ["CCM"], per_family=1)) == 1
|
|
assert len(manual_sources(cpt_replica, ["CCM"], per_family=2)) == 2
|
|
|
|
def test_snippet_chars_bounds_the_text_passed_to_as_source(self, cpt_replica):
|
|
out = manual_sources(cpt_replica, ["CCM"], snippet_chars=10)
|
|
assert len(out[0]["snippet"]) <= 10
|
|
|
|
def test_missing_tables_yield_empty(self):
|
|
con = duckdb.connect(":memory:")
|
|
try:
|
|
assert manual_sources(con, ["CCM"]) == []
|
|
finally:
|
|
con.close()
|