feat(llm): lineage evidence — timeline events, element diffs and guidance as a lineage SSE event + cited prompt block (refs #691)

This commit is contained in:
kert
2026-09-10 10:26:34 -04:00
parent 78da985b3a
commit 6924b9b6c3
8 changed files with 1562 additions and 24 deletions

View File

@@ -49,6 +49,8 @@ class LlmConfig:
code_cited_per_code: int = 3
code_cited_collections: tuple[str, ...] = ("rules", "comments", "corpus")
code_cited_max: int = 12
lineage_max_rows: int = 25
lineage_on_demand_max: int = 3
def parse_hosts(spec: str) -> tuple[tuple[str, float], ...]:
@@ -120,6 +122,8 @@ def load() -> LlmConfig:
_opt(section, "code_cited_collections", ["rules", "comments", "corpus"])
),
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)),
embed_dim=int(section.embed_dim),
build_ann_index=build_ann_index,
pg_host=os.environ.get("LLM_PG_HOST", str(section.pg_host)),

View File

@@ -393,3 +393,10 @@ def code_cited_sources(
def merge_sources(retrieved: list[dict], extra: list[dict]) -> list[dict]:
labels = {s["label"] for s in retrieved}
return retrieved + [s for s in extra if s["label"] not in labels]
# Re-exported so rag.py imports both evidence functions from one place,
# like valuation_evidence — llm.lineage lazily imports this module's
# _connect/_replica_path/_mtime/warm inside lineage_evidence() rather
# than at module scope, so this import carries no circularity risk.
from llm.lineage import lineage_evidence # noqa: E402,F401

510
src/llm/lineage.py Normal file
View File

@@ -0,0 +1,510 @@
"""Structured lineage evidence for the chat: a dated timeline, cross-code
element differences, and CFR/IOM/MLN guidance references, beside the
excerpts and the valuation table.
``lineage_evidence`` detects HCPCS/CPT codes (and registered families)
in the question, then reads the precomputed ``pfs.code_event`` rows
(``pfs.codetables.read_events`` — every A/R/T code, Task 2's
``stack pfs lineage --all-payable --write``) from the cached DuckDB
read replica. A code with no precomputed rows falls back on demand to
``pfs.lineage.lineage`` (bounded by ``cfg.lineage_on_demand_max`` codes
per turn — that path re-derives events from the RVU files and a live
scan of ``fr_anchors``, too slow to run for every code on every turn).
Events for the same ``(code, year, kind, from_codes, to_codes)`` are
collapsed to one representative row, preferring an anchored FR
paragraph. ``pfs.code_element`` supplies element differences across the
detected codes' newest year, and ``pfs.code_guidance`` supplies CFR/IOM/
MLN references for the detected families. Like ``llm.evidence``, this
degrades to "no evidence" on any failure — the chat never fails because
of it.
"""
from __future__ import annotations
import logging
import threading
from dataclasses import asdict, dataclass
from typing import Any, Sequence
from llm.config import LlmConfig
from pfs.codetables import (
ElementRow,
EventRow,
GuidanceRow,
is_missing_table_error,
read_elements,
read_events,
read_guidance,
)
from pfs.descriptors import rule_year_of
from pfs.families import detect_codes
log = logging.getLogger(__name__)
_HEADER = (
"Lineage (dated events with anchors; cite the bracketed label for any dated claim):"
)
#: Event kinds a prompt block never drops for the row cap (ruling 17) —
#: these are the "what happened and what replaced it" facts a dated
#: claim usually hinges on; revaluation/descriptor noise can be trimmed
#: first.
_PRIORITY_KINDS = frozenset(
{"created", "adopted_cpt", "replaces", "replaced_by", "deleted", "disappeared"}
)
#: Sort order for collapsed events sharing a year and code: FR event
#: kinds (their own internal order), then RVU-file kinds, then the two
#: CPT-codebook kinds. Anything unrecognized sorts last.
_KIND_ORDER = (
"created",
"adopted_cpt",
"replaces",
"replaced_by",
"deleted",
"crosswalk",
"bundled",
"telehealth_list",
"appeared",
"disappeared",
"status_change",
"descriptor_change",
"revalued",
"cpt_changed",
"cpt_deleted",
)
#: Guidance rows kept in the prompt/payload per question — a code's
#: CFR/IOM/MLN references rarely run past a handful, and a long tail
#: would crowd the lineage block out of the prompt budget.
_GUIDANCE_CAP = 8
#: Element diffs kept, most-shared value first (ruling 18).
_ELEMENT_DIFF_CAP = 12
def _kind_rank(kind: str) -> int:
try:
return _KIND_ORDER.index(kind)
except ValueError:
return len(_KIND_ORDER)
@dataclass(frozen=True)
class LineageEvent:
code: str
year: int
kind: str
from_codes: tuple[str, ...]
to_codes: tuple[str, ...]
label: str
item_key: str
p_id: int
page: int
url: str
source: str # "fr" | "rvu" | "cpt"
anchored: bool
note: str
@dataclass(frozen=True)
class ElementDiff:
type: str
value: str
in_codes: tuple[str, ...]
not_in_codes: tuple[str, ...]
label: str
item_key: str
p_id: int
@dataclass(frozen=True)
class GuidanceRef:
kind: str # "cfr" | "iom" | "mln"
locator: str
url: str
label: str
item_key_src: str
p_id_src: int
def _event_line(e: LineageEvent) -> str:
line = f"[{e.label}] {e.year} {e.kind} {e.code}"
if e.from_codes or e.to_codes:
line += f" ({', '.join(e.from_codes)}{', '.join(e.to_codes)})"
if e.note:
line += f"{e.note}"
return line
def _diff_line(d: ElementDiff) -> str:
return (
f"[{d.label}] {d.type}={d.value}: "
f"in {', '.join(d.in_codes)}; not in {', '.join(d.not_in_codes)}"
)
def _guidance_line(g: GuidanceRef) -> str:
return f"[{g.label}] {g.locator}{g.kind.upper()}"
@dataclass(frozen=True)
class LineageEvidence:
codes: tuple[str, ...]
families: tuple[str, ...]
#: Every collapsed event row, sorted (year, code, kind order) — the
#: SSE payload carries all of these; ``prompt_block`` trims them.
events: tuple[LineageEvent, ...]
element_diffs: tuple[ElementDiff, ...]
guidance: tuple[GuidanceRef, ...]
#: "elements extracted for N of M codes" when only one detected code
#: has any ``pfs.code_element`` rows (ruling 18) — "" otherwise.
elements_note: str = ""
#: ``cfg.lineage_max_rows`` at construction time — ``prompt_block``
#: is self-contained (no cfg argument), like ``ValuationEvidence``.
max_prompt_rows: int = 25
def _prompt_events(self) -> list[LineageEvent]:
priority = [e for e in self.events if e.kind in _PRIORITY_KINDS]
rest = [e for e in self.events if e.kind not in _PRIORITY_KINDS]
budget = max(self.max_prompt_rows - len(priority), 0)
keep = set(priority) | set(rest[:budget])
return [e for e in self.events if e in keep]
def prompt_block(self) -> str:
lines = [_HEADER]
lines.extend(_event_line(e) for e in self._prompt_events())
if self.element_diffs:
lines.append("Element differences:")
lines.extend(_diff_line(d) for d in self.element_diffs)
if self.guidance:
lines.append("Guidance:")
lines.extend(_guidance_line(g) for g in self.guidance)
return "\n".join(lines)
def payload(self) -> dict[str, Any]:
out: dict[str, Any] = {
"type": "lineage",
"codes": list(self.codes),
"families": list(self.families),
"events": [asdict(e) for e in self.events],
"element_diffs": [asdict(d) for d in self.element_diffs],
"guidance": [asdict(g) for g in self.guidance],
}
if self.elements_note:
out["elements_note"] = self.elements_note
return out
def rule_label(title: str, date_published: str, p_id: int) -> str:
"""``"CY2021 PFS final ¶1578"`` / ``"CY2027 PFS proposed ¶394"`` —
``proposed`` when *title* contains "Proposed" (case-insensitive),
else ``final``; the year from ``rule_year_of``."""
year = rule_year_of(title, date_published)
kind = "proposed" if "proposed" in (title or "").lower() else "final"
return f"CY{year} PFS {kind}{p_id}"
# ── The bib Store, opened lazily once per process (sqlite, cheap) ──────
_STORE: Any | None = None
_STORE_LOCK = threading.Lock()
def _store() -> Any:
global _STORE
with _STORE_LOCK:
if _STORE is None:
from conf.connect import bib
_STORE = bib()
return _STORE
def _source_label(store: Any, item_key: str, p_id: int) -> str:
"""``rule_label`` for the FR rule at *item_key* — falls back to the
bare item key when the item can't be fetched (unresolvable/missing
item must never break the chat over a label)."""
try:
item = store.get(item_key)
except Exception as e: # noqa: BLE001
log.warning("lineage label unresolved (%s): %s", item_key, e)
return item_key or "?"
return rule_label(item.title, item.date_published, p_id)
# ── FR jump links, cached per (item_key, p_id) ──────────────────────────
_URL_CACHE: dict[tuple[str, int], str] = {}
_URL_LOCK = threading.Lock()
def _fr_url(store: Any, item_key: str, p_id: int) -> str:
if not item_key or not p_id:
return ""
key = (item_key, p_id)
with _URL_LOCK:
cached = _URL_CACHE.get(key)
if cached is not None:
return cached
url = ""
try:
from bib import frlink
url = frlink.resolve(f"p-{p_id}", store=store, item_key=item_key).url
except Exception as e: # noqa: BLE001 — a bad anchor must not break the chat
log.warning("lineage FR url unresolved (%s p-%s): %s", item_key, p_id, e)
with _URL_LOCK:
_URL_CACHE[key] = url
return url
def _to_lineage_event(store: Any, r: EventRow) -> LineageEvent:
from_codes = tuple(r.from_codes.split()) if r.from_codes else ()
to_codes = tuple(r.to_codes.split()) if r.to_codes else ()
if r.source == "fr":
label = _source_label(store, r.item_key, r.p_id)
url = _fr_url(store, r.item_key, r.p_id)
elif r.source == "cpt":
label = r.note or f"CPT Changes {r.year}"
url = ""
else: # "rvu"
label = f"PFS CY{r.year} RVU file"
url = ""
return LineageEvent(
code=r.code,
year=r.year,
kind=r.kind,
from_codes=from_codes,
to_codes=to_codes,
label=label,
item_key=r.item_key,
p_id=r.p_id,
page=r.page,
url=url,
source=r.source,
anchored=r.anchored,
note=r.note,
)
def _collapse(rows: Sequence[EventRow]) -> list[EventRow]:
"""One representative row per ``(code, year, kind, from_codes,
to_codes)`` — prefer an anchored FR row, then the lowest ``p_id`` —
sorted by ``(year, code, kind order)``."""
groups: dict[tuple[str, int, str, str, str], list[EventRow]] = {}
for r in rows:
key = (r.code, r.year, r.kind, r.from_codes, r.to_codes)
groups.setdefault(key, []).append(r)
reps = [
min(
grp,
key=lambda r: (0 if (r.source == "fr" and r.anchored) else 1, r.p_id),
)
for grp in groups.values()
]
return sorted(reps, key=lambda r: (r.year, r.code, _kind_rank(r.kind)))
# ── On-demand fallback for codes with no precomputed rows, cached by
# (replica mtime, code) — see the module docstring. ─────────────────────
_ON_DEMAND_CACHE: dict[str, list[EventRow]] = {}
_ON_DEMAND_MTIME: int | None = None
_ON_DEMAND_LOCK = threading.Lock()
def _on_demand_events(cur: Any, store: Any, mtime: int, code: str) -> list[EventRow]:
global _ON_DEMAND_MTIME
with _ON_DEMAND_LOCK:
if _ON_DEMAND_MTIME != mtime:
_ON_DEMAND_CACHE.clear()
_ON_DEMAND_MTIME = mtime
cached = _ON_DEMAND_CACHE.get(code)
if cached is not None:
return cached
try:
import pfs.lineage as pfs_lineage
rows = pfs_lineage.lineage(cur, store, code)
except Exception as e: # noqa: BLE001 — on-demand fallback must not break the chat
log.warning("on-demand lineage skipped for %s: %s", code, e)
rows = []
with _ON_DEMAND_LOCK:
if _ON_DEMAND_MTIME == mtime:
_ON_DEMAND_CACHE[code] = rows
return rows
def _collect_events(
cur: Any, store: Any, mtime: int, codes: Sequence[str], on_demand_max: int
) -> list[EventRow]:
out: list[EventRow] = []
used = 0
for code in codes:
try:
rows = read_events(cur, code)
except Exception as e: # noqa: BLE001
if not is_missing_table_error(e):
raise
rows = []
if not rows and used < on_demand_max:
rows = _on_demand_events(cur, store, mtime, code)
used += 1
out.extend(rows)
return out
def _element_label(store: Any, row: ElementRow) -> str:
if row.item_key and row.p_id:
return _source_label(store, row.item_key, row.p_id)
if row.source == "cpt":
return f"CPT Changes {row.year}"
return f"PFS CY{row.year} RVU file"
def _element_diffs(
cur: Any, codes: Sequence[str], store: Any
) -> tuple[tuple[ElementDiff, ...], str]:
"""Element values present for some detected codes and absent for
others, from each code's newest ``pfs.code_element`` year. Needs
>= 2 detected codes with any element rows at all — a single code's
elements have nothing to diff against (ruling 18)."""
by_code: dict[str, list[ElementRow]] = {}
for code in codes:
try:
rows = read_elements(cur, code)
except Exception as e: # noqa: BLE001
if not is_missing_table_error(e):
raise
rows = []
if not rows:
continue
newest = max(r.year for r in rows)
by_code[code] = [r for r in rows if r.year == newest]
n, m = len(by_code), len(codes)
if n < 2:
note = f"elements extracted for {n} of {m} codes" if n == 1 else ""
return (), note
groups: dict[tuple[str, str], list[tuple[str, ElementRow]]] = {}
for code, rows in by_code.items():
seen: set[tuple[str, str]] = set()
for r in rows:
pair = (r.type, r.value)
if pair in seen:
continue
seen.add(pair)
groups.setdefault(pair, []).append((code, r))
all_with_elements = set(by_code)
diffs: list[ElementDiff] = []
for (type_, value_), entries in groups.items():
in_codes = tuple(sorted({code for code, _ in entries}))
not_in_codes = tuple(sorted(all_with_elements - set(in_codes)))
if not not_in_codes:
continue # present for every code that has elements at all
best = min(
(row for _, row in entries),
key=lambda r: (0 if (r.item_key and r.p_id) else 1, r.year),
)
diffs.append(
ElementDiff(
type=type_,
value=value_,
in_codes=in_codes,
not_in_codes=not_in_codes,
label=_element_label(store, best),
item_key=best.item_key,
p_id=best.p_id,
)
)
diffs.sort(key=lambda d: (-len(d.in_codes), d.type, d.value))
return tuple(diffs[:_ELEMENT_DIFF_CAP]), ""
def _collect_guidance(
cur: Any, store: Any, families: Sequence[str], wide: Sequence[str]
) -> tuple[GuidanceRef, ...]:
wide_set = set(wide)
rows: list[GuidanceRow] = []
for family in families:
if family in wide_set:
continue
try:
rows.extend(read_guidance(cur, family))
except Exception as e: # noqa: BLE001
if not is_missing_table_error(e):
raise
rows.sort(key=lambda r: (0 if r.kind == "cfr" else 1, r.kind, r.locator))
out: list[GuidanceRef] = []
seen: set[str] = set()
for r in rows:
if r.locator in seen:
continue
seen.add(r.locator)
url = ""
if r.kind == "cfr":
try:
from bib import cfrlink
url = cfrlink.url(cfrlink.parse_cite(r.locator))
except Exception as e: # noqa: BLE001
log.warning("lineage CFR url unresolved (%s): %s", r.locator, e)
out.append(
GuidanceRef(
kind=r.kind,
locator=r.locator,
url=url,
label=_source_label(store, r.item_key_src, r.p_id_src),
item_key_src=r.item_key_src,
p_id_src=r.p_id_src,
)
)
if len(out) >= _GUIDANCE_CAP:
break
return tuple(out)
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
raises: any failure opening the replica or querying it is logged and
degrades to ``None``, the same contract as ``valuation_evidence``."""
import llm.evidence as evidence
evidence.warm(cfg)
det = detect_codes(question)
if not det.codes:
return None
path = evidence._replica_path(cfg)
try:
con = evidence._connect(path)
except Exception as e: # noqa: BLE001 — duckdb raises several types
log.warning("lineage evidence skipped (%s): %s", path, e)
return None
cur = con.cursor()
try:
store = _store()
mtime = evidence._mtime(path)
raw_events = _collect_events(
cur, store, mtime, det.codes, cfg.lineage_on_demand_max
)
events = tuple(_to_lineage_event(store, r) for r in _collapse(raw_events))
element_diffs, elements_note = _element_diffs(cur, det.codes, store)
guidance = _collect_guidance(cur, store, det.families, det.wide)
except Exception as e: # noqa: BLE001
log.warning("lineage evidence skipped (query): %s", e)
return None
finally:
cur.close() # the cached parent connection stays open
return LineageEvidence(
codes=det.codes,
families=det.families,
events=events,
element_diffs=element_diffs,
guidance=guidance,
elements_note=elements_note,
max_prompt_rows=cfg.lineage_max_rows,
)

View File

@@ -8,10 +8,11 @@ merged and re-ranked by similarity × recency (``llm.rerank``), then each
source gets a deep link (``llm.links``). Generation streams from the
largest live Ollama host (``HostPool.acquire_generation``) with the
model tier that host can hold (``pick_model``). When the question
cites HCPCS/CPT codes, ``llm.evidence`` supplies authoritative
valuation rows and code- or family-cited chunks from rules, comments and
the reference corpus; the SSE event order is ``valuation`` (only when
evidence is found) → ``token``* → ``sources`` → ``done``.
cites HCPCS/CPT codes, ``llm.evidence``/``llm.lineage`` supply
authoritative valuation rows, a dated lineage timeline, and code- or
family-cited chunks from rules, comments and the reference corpus; the
SSE event order is ``lineage`` (only when detected) → ``valuation``
(only when found) → ``token``* → ``sources`` → ``done``.
"""
from __future__ import annotations
@@ -26,10 +27,12 @@ from llm.config import LlmConfig
from llm.evidence import (
ValuationEvidence,
code_cited_sources,
lineage_evidence,
merge_sources,
valuation_evidence,
)
from llm.index import _engine
from llm.lineage import LineageEvidence
from llm.links import as_source
from llm.pool import HostPool, PoolEmbeddings, pick_model
from llm.rerank import Hit, blend, filter_since
@@ -57,6 +60,11 @@ _SYSTEM = (
"e.g. [PFS CY2026 Addendum B]. Do not compute, extrapolate or convert "
"numbers beyond what the rows show; if a code is listed as not priced, "
"say so."
" A Lineage section may also follow. It is authoritative for when a "
"code was created, adopted, replaced or deleted and what replaced it; "
"cite its bracketed label for every dated claim. If the Lineage "
"section does not give a date for something, say so rather than "
"inventing one."
)
@@ -109,11 +117,14 @@ def retrieve(
def build_messages(
question: str, sources: list[dict], evidence: ValuationEvidence | None = None
question: str,
sources: list[dict],
evidence: ValuationEvidence | None = None,
lineage: LineageEvidence | None = None,
) -> list[dict]:
"""Grounded chat messages: system rules + question with excerpts and,
when available, a Valuation block between the excerpts and the
question."""
when available, a Lineage block and a Valuation block — in that
order — between the excerpts and the question."""
if sources:
context = "\n\n".join(
f"[{s['label']}] ({s.get('kind', '')}, {s.get('date', '') or 'undated'}) "
@@ -123,6 +134,8 @@ def build_messages(
else:
context = "(no relevant excerpts found)"
parts = [f"Excerpts:\n\n{context}"]
if lineage is not None:
parts.append(lineage.prompt_block())
if evidence is not None:
parts.append(evidence.prompt_block())
parts.append(f"Question: {question}")
@@ -138,30 +151,39 @@ def stream_answer(
) -> Iterator[dict]:
"""Retrieve, then stream a grounded answer from the largest live host.
When the question cites HCPCS/CPT codes, yields a ``valuation`` event
first (evidence's ``payload()``) and merges chunks from rules,
comments and the corpus that cite those codes — or, for a detected
family, its family key — into ``sources``. Then yields
``{"type":"token","text":…}``
events as the model generates, then one
``{"type":"sources", "sources": […], "model": …, "host": …}`` and a
final ``{"type":"done"}``.
When the question cites HCPCS/CPT codes, yields a ``lineage`` event
(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
yields ``{"type":"token","text":…}`` events as the model generates,
then one ``{"type":"sources", "sources": […], "model": …, "host":
…}`` and a final ``{"type":"done"}``.
"""
sources = retrieve(question, cfg=cfg, pool=pool, since=since)
lineage = lineage_evidence(question, cfg)
evidence = valuation_evidence(question, cfg)
if evidence is not None:
codes: tuple[str, ...] = evidence.codes if evidence is not None else ()
families: tuple[str, ...] = evidence.families if evidence is not None else ()
if lineage is not None:
codes = tuple(sorted(set(codes) | set(lineage.codes)))
families = tuple(sorted(set(families) | set(lineage.families)))
if evidence is not None or lineage is not None:
cited = code_cited_sources(
_engine(cfg),
evidence.codes,
codes,
per_code=cfg.code_cited_per_code,
collections=tuple(cfg.code_cited_collections),
families=evidence.families,
families=families,
max_total=cfg.code_cited_max,
)
sources = merge_sources(sources, cited)
if lineage is not None:
yield lineage.payload()
if evidence is not None:
yield evidence.payload()
pool.check(cfg.instruct_model)
messages = build_messages(question, sources, evidence)
messages = build_messages(question, sources, evidence, lineage)
with pool.acquire_generation() as host, httpx.Client(timeout=_TIMEOUT) as client:
model = pick_model(cfg, pool, host)
with client.stream(

View File

@@ -129,6 +129,8 @@ valuation_years = 4 # final-rule vintages shown per code (p
code_cited_per_code = 3 # excerpts literally citing each detected code (or family)
code_cited_collections = ["rules", "comments", "corpus"] # searched in this order
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
[llm.k_per_kind] # over-fetched ×3 per kind, then re-ranked
comment = 8

View File

@@ -96,6 +96,8 @@ class TestNewKnobs:
assert cfg.code_cited_per_code == 3
assert cfg.code_cited_collections == ("rules", "comments", "corpus")
assert cfg.code_cited_max == 12
assert cfg.lineage_max_rows == 25
assert cfg.lineage_on_demand_max == 3
def test_duckdb_replica_env_override(self, monkeypatch):
monkeypatch.setenv("LLM_DUCKDB_REPLICA", "/app/data/replica/aco.ro.duckdb")

775
tests/llm/test_lineage.py Normal file
View File

@@ -0,0 +1,775 @@
"""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,
lineage_evidence,
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._URL_CACHE.clear()
lineage._ON_DEMAND_CACHE.clear()
lineage._ON_DEMAND_MTIME = None
_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);"
"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",
),
(
"YBM4IZUS",
"Medicare Program; CY 2021 Payment Policies Under the PFS",
"2020-12-28",
),
(
"ZPROP2027",
"Medicare Program; CY 2027 Proposed Payment Policies",
"2026-07-16",
),
],
)
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
class TestElementDiffs:
def test_no_diffs_when_fewer_than_two_codes_have_elements(self, con):
c, _ = con
assert lineage._element_diffs(c, ["99490"], None) == ((), "")
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)
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)
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)
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)
# 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)
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"], [])
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"], [])
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"])
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"], [])
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 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}"

View File

@@ -1,5 +1,6 @@
"""llm.rag — multi-collection retrieval + grounded streaming answer."""
import json
from datetime import date
from unittest.mock import MagicMock, patch
@@ -7,6 +8,7 @@ import pytest
from langchain_core.documents import Document
from llm.config import LlmConfig
from llm.lineage import LineageEvent, LineageEvidence
from llm.rag import build_messages, retrieve, stream_answer
CFG = LlmConfig(
@@ -287,6 +289,57 @@ class TestBuildMessages:
assert build_messages("q", []) == build_messages("q", [], evidence=None)
assert "Valuation" not in build_messages("q", [])[1]["content"]
def test_lineage_block_between_excerpts_and_valuation(self):
from llm.evidence import ValuationEvidence
le = LineageEvidence(
codes=("G2058",),
families=("CCM",),
events=(
LineageEvent(
"G2058",
2021,
"replaced_by",
(),
("99439",),
"CY2021 PFS final ¶686",
"YBM4IZUS",
686,
0,
"u",
"fr",
True,
"",
),
),
element_diffs=(),
guidance=(),
)
ev = ValuationEvidence(("G2058",), ("CCM",), (), ())
msgs = build_messages("history?", [], evidence=ev, lineage=le)
user = msgs[1]["content"]
assert (
user.index("Excerpts:")
< user.index("Lineage (dated events")
< user.index("Valuation (authoritative")
< user.index("Question: history?")
)
assert "[CY2021 PFS final ¶686] 2021 replaced_by G2058" in user
assert "lineage" in msgs[0]["content"].lower()
assert "cite its bracketed label for every dated claim" in msgs[0]["content"]
def test_lineage_only_no_valuation_block(self):
le = LineageEvidence(
codes=("G2058",), families=(), events=(), element_diffs=(), guidance=()
)
msgs = build_messages("q", [], lineage=le)
user = msgs[1]["content"]
assert "Lineage (dated events" in user
assert "Valuation (authoritative" not in user
def test_no_lineage_prompt_unchanged(self):
assert build_messages("q", []) == build_messages("q", [], lineage=None)
class TestStreamAnswer:
def _pool(self, vram=24.0, serves_big=True):
@@ -298,10 +351,11 @@ class TestStreamAnswer:
@patch("llm.rag._engine")
@patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.lineage_evidence", return_value=None)
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_yields_tokens_then_sources_then_done(
self, mock_retrieve, MockClient, _ev, mock_engine
self, mock_retrieve, MockClient, _lin, _ev, mock_engine
):
src = {
"id": "C1",
@@ -343,9 +397,10 @@ class TestStreamAnswer:
assert client.stream.call_args.args[1] == "http://h1:11434/api/chat"
@patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.lineage_evidence", return_value=None)
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_small_host_uses_baseline_model(self, mock_retrieve, MockClient, _ev):
def test_small_host_uses_baseline_model(self, mock_retrieve, MockClient, _lin, _ev):
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
@@ -355,9 +410,10 @@ class TestStreamAnswer:
assert events[-2]["model"] == "chat"
@patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.lineage_evidence", return_value=None)
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_since_forwarded(self, mock_retrieve, MockClient, _ev):
def test_since_forwarded(self, mock_retrieve, MockClient, _lin, _ev):
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
@@ -366,9 +422,10 @@ class TestStreamAnswer:
assert mock_retrieve.call_args.kwargs["since"] == "2025-09-01"
@patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.lineage_evidence", return_value=None)
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_http_error_propagates(self, mock_retrieve, MockClient, _ev):
def test_http_error_propagates(self, mock_retrieve, MockClient, _lin, _ev):
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
@@ -379,10 +436,11 @@ class TestStreamAnswer:
@patch("llm.rag._engine")
@patch("llm.rag.code_cited_sources")
@patch("llm.rag.valuation_evidence")
@patch("llm.rag.lineage_evidence", return_value=None)
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_valuation_event_first_and_sources_merged(
self, mock_retrieve, MockClient, mock_ev, mock_cited, mock_engine
self, mock_retrieve, MockClient, _lin, mock_ev, mock_cited, mock_engine
):
from llm.evidence import ValuationEvidence
@@ -422,3 +480,161 @@ class TestStreamAnswer:
)
body = client.stream.call_args.kwargs["json"]
assert "Valuation (authoritative" in body["messages"][1]["content"]
@patch("llm.rag._engine")
@patch("llm.rag.code_cited_sources")
@patch("llm.rag.valuation_evidence")
@patch("llm.rag.lineage_evidence")
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_lineage_event_before_valuation_and_tokens(
self, mock_retrieve, MockClient, mock_lin, mock_ev, mock_cited, mock_engine
):
from llm.evidence import ValuationEvidence
mock_retrieve.return_value = []
mock_cited.return_value = []
mock_lin.return_value = LineageEvidence(
codes=("G2058",),
families=("CCM",),
events=(
LineageEvent(
"G2058",
2021,
"replaced_by",
(),
("99439",),
"CY2021 PFS final ¶686",
"YBM4IZUS",
686,
0,
"u",
"fr",
True,
"",
),
),
element_diffs=(),
guidance=(),
)
mock_ev.return_value = ValuationEvidence(("G2058",), ("CCM",), (), ())
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
events = list(stream_answer("history of G2058?", cfg=CFG, pool=self._pool()))
assert [e["type"] for e in events] == [
"lineage",
"valuation",
"token",
"sources",
"done",
]
assert events[0]["codes"] == ["G2058"]
body = client.stream.call_args.kwargs["json"]
content = body["messages"][1]["content"]
assert content.index("Lineage (dated events") < content.index(
"Valuation (authoritative"
)
@patch("llm.rag._engine")
@patch("llm.rag.code_cited_sources")
@patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.lineage_evidence")
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_lineage_only_calls_cited_sources_with_lineage_codes(
self, mock_retrieve, MockClient, mock_lin, _ev, mock_cited, mock_engine
):
mock_retrieve.return_value = []
mock_cited.return_value = []
mock_lin.return_value = LineageEvidence(
codes=("99490", "99491"),
families=("CCM",),
events=(),
element_diffs=(),
guidance=(),
)
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
events = list(stream_answer("CCM history?", cfg=CFG, pool=self._pool()))
assert events[0]["type"] == "lineage"
mock_cited.assert_called_once_with(
mock_engine.return_value,
("99490", "99491"),
per_code=CFG.code_cited_per_code,
collections=CFG.code_cited_collections,
families=("CCM",),
max_total=CFG.code_cited_max,
)
@patch("llm.rag._engine")
@patch("llm.rag.code_cited_sources")
@patch("llm.rag.valuation_evidence")
@patch("llm.rag.lineage_evidence")
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_codes_and_families_unioned_when_both_present(
self, mock_retrieve, MockClient, mock_lin, mock_ev, mock_cited, mock_engine
):
from llm.evidence import ValuationEvidence
mock_retrieve.return_value = []
mock_cited.return_value = []
# lineage reaches a code (G2058) with no RVU rows valuation never sees.
mock_lin.return_value = LineageEvidence(
codes=("99490", "G2058"),
families=("CCM",),
events=(),
element_diffs=(),
guidance=(),
)
mock_ev.return_value = ValuationEvidence(("99490",), ("CCM",), (), ())
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
list(stream_answer("CCM and G2058?", cfg=CFG, pool=self._pool()))
mock_cited.assert_called_once_with(
mock_engine.return_value,
("99490", "G2058"),
per_code=CFG.code_cited_per_code,
collections=CFG.code_cited_collections,
families=("CCM",),
max_total=CFG.code_cited_max,
)
@patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.httpx.Client")
@patch("llm.rag.retrieve")
def test_control_question_events_match_the_no_lineage_baseline(
self, mock_retrieve, MockClient, _ev
):
"""No codes in the question — the real lineage_evidence short-
circuits on an empty Detection without touching the replica, so
its event stream must be byte-identical to one where the feature
is switched off outright (``lineage_evidence`` patched to
``None``)."""
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
def _run():
resp.iter_lines.return_value = iter(
['{"message":{"content":"x"},"done":true}']
)
return list(
stream_answer(
"why did CMS finalize this policy?", cfg=CFG, pool=self._pool()
)
)
with_real_lineage = json.dumps(_run())
with patch("llm.rag.lineage_evidence", return_value=None):
with_lineage_off = json.dumps(_run())
assert with_real_lineage == with_lineage_off