diff --git a/src/llm/config.py b/src/llm/config.py index 63584e6..aa88a71 100644 --- a/src/llm/config.py +++ b/src/llm/config.py @@ -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)), diff --git a/src/llm/evidence.py b/src/llm/evidence.py index de03e2f..ed10291 100644 --- a/src/llm/evidence.py +++ b/src/llm/evidence.py @@ -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 diff --git a/src/llm/lineage.py b/src/llm/lineage.py new file mode 100644 index 0000000..6584a0b --- /dev/null +++ b/src/llm/lineage.py @@ -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, + ) diff --git a/src/llm/rag.py b/src/llm/rag.py index b24fc4b..b003f31 100644 --- a/src/llm/rag.py +++ b/src/llm/rag.py @@ -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( diff --git a/stack.toml b/stack.toml index 172d737..76970d2 100644 --- a/stack.toml +++ b/stack.toml @@ -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 diff --git a/tests/llm/test_config.py b/tests/llm/test_config.py index ca394cd..6b6b305 100644 --- a/tests/llm/test_config.py +++ b/tests/llm/test_config.py @@ -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") diff --git a/tests/llm/test_lineage.py b/tests/llm/test_lineage.py new file mode 100644 index 0000000..1fe8918 --- /dev/null +++ b/tests/llm/test_lineage.py @@ -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}" diff --git a/tests/llm/test_rag.py b/tests/llm/test_rag.py index 1859f9a..f976466 100644 --- a/tests/llm/test_rag.py +++ b/tests/llm/test_rag.py @@ -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