Files
stack/tests/llm/test_evidence.py

234 lines
7.6 KiB
Python

"""llm.evidence — valuation prompt block, SSE payload, code-cited rule sources."""
from __future__ import annotations
from unittest.mock import MagicMock, patch
from llm.config import LlmConfig
from llm.evidence import (
ValuationEvidence,
code_cited_sources,
merge_sources,
valuation_evidence,
)
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,
)
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_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["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_called_once()
@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
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,
}
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
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_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 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