homelab/stack
1
0
Fork 0
Code Issues 114 Pull Requests 1 Actions Packages Projects 1 Releases Wiki Activity

fix(llm): thread-local bib Store; deterministic collapse; unique labels; era coverage; system prompt; prompt budget (refs #691 #692)

Browse Source 
Final fix-wave items C1, I1-I4 and B11 (refined), all landing in the
same handful of interconnected files (llm.lineage/llm.rag/llm.evidence
share the collapse/select/label/budget call paths, so they can't be
split into independently-working commits):

- C1 (critical): llm.lineage._store() was one process-global bib.Store
  whose sqlite connection is check_same_thread=True — /chat runs each
  turn in Starlette's threadpool, so every thread but the first got a
  silently inert (or raising) store. Now threading.local(), one Store
  opened lazily per thread. rag._docket_year uses the same accessor and
  is now fully guarded (never escapes era_of). New
  TestConcurrentLineage (8 threads x 40 calls, mirrors
  TestConcurrentChats) asserts no cross-thread error and one Store
  construction per thread.

- I1 (Ruling B12): _collapse's tie-break is now
  (not anchored_fr, is_proposed, p_id, item_key) — a final rule beats a
  tied proposed one, and a fully-tied pair is decided by item_key for a
  deterministic total order.

- I2 (Ruling B13): rule_label carries "{vol} FR {page}" when the
  paragraph resolves (via the cached item/paragraph lookups), making
  labels unique across paragraphs that used to share one.
  merge_sources dedupes rule-kind rows on (item_key, p_id) instead of
  label, since two rule chunks for the same paragraph can now carry
  different labels.

- I3 (Important): era_balance picks eras evenly across the range when
  there are more distinct eras than top_n, so the round-robin (which
  visits newest-first every round) doesn't silently drop the oldest
  eras from a wide history question.

- I4 (Important): _SYSTEM's opener, refusal clause and recency
  guidance are reworded to match what the prompt actually contains
  (excerpts + optional Lineage + optional Valuation), and the Lineage
  paragraph now precedes Valuation to match build_messages' order.

- B11 refined (whole-branch review): a hard per-turn code cap
  (chat_codes_max=24, explicit codes then family order —
  llm.evidence.cap_codes, shared by valuation_evidence and
  lineage_evidence so both cap the same question identically),
  a valuation-rows cap (valuation_rows_max=24, explicit then newest
  vintage), a lineage-rows hard cap (2x lineage_max_rows, priority
  rows capped at 4/code), element-diff code lists compacted past 8,
  manual sources capped at 2, and build_messages(budget_chars=...)
  which drops lineage-source excerpts>4, retrieved>6, cited>8,
  manual>1, valuation rows>12, then lineage rows>lineage_max_rows in
  that order until the assembled prompt fits — wired into
  stream_answer as budget_chars=cfg.chat_num_ctx*3. The
  _LINEAGE_SOURCES_HARD_CAP is now enforced in lineage_sources' own
  event loop, not only its element-diff tail.

I5: TestLineageEvidenceLive now uses the shared restore_families
fixture (moved to tests/conftest.py) so opening the real replica
doesn't leak thousands of derived families into later tests.
Diagnosing this also turned up a second, pre-existing leak of the same
shape: TestStreamAnswer's control-question test ran the real
lineage_evidence against CFG's default (real, 3GB)
data/replica/aco.ro.duckdb, since llm.lineage.lineage_evidence calls
evidence.warm(cfg) unconditionally before checking for detected codes
— fixed by pointing that one test at a nonexistent replica path
(exactly the "without touching the replica" behavior its own docstring
already claimed).

Verified: uv run pytest tests/llm tests/pfs/test_families.py
tests/pfs/test_lineage.py tests/cli/test_pfs_cli.py tests/dev -q
-p no:cacheprovider -m "not live" — 583 passed. In-process golden run
against the live replica: 3/7 pass (ccm-history, audio-only-em-99441,
g2211-commenters); g2058-replacement/apcm-vs-ccm/99490-telehealth
unchanged documented gaps; g2064-g2065 newly misses one of its three
anchors (JE7KYBW3 p1111) specifically because of the new 4-per-code
lineage-row cap this commit adds (Ruling B11) — an accepted tradeoff
of the budget work, not a bug. Prompt-size check (budget 24,576
chars): "history of CCM coding and payment" 20,670 chars; the 58-code
three-family history question 16,135 chars (58 detected codes capped
to 24) — both under budget.
This commit is contained in:
kert
2026-09-10 13:19:03 -04:00
parent a5a5f60594
commit c1a199c878
10 changed files with 1243 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/llm/config.py
View File

@@ -54,6 +54,8 @@ class LlmConfig:
lineage_max_rows: int = 25
lineage_on_demand_max: int = 3
lineage_sources_max: int = 8
chat_codes_max: int = 24
valuation_rows_max: int = 24
def parse_hosts(spec: str) -> tuple[tuple[str, float], ...]:
@@ -130,6 +132,8 @@ def load() -> LlmConfig:
lineage_max_rows=int(_opt(section, "lineage_max_rows", 25)),
lineage_on_demand_max=int(_opt(section, "lineage_on_demand_max", 3)),
lineage_sources_max=int(_opt(section, "lineage_sources_max", 8)),
chat_codes_max=int(_opt(section, "chat_codes_max", 24)),
valuation_rows_max=int(_opt(section, "valuation_rows_max", 24)),
embed_dim=int(section.embed_dim),
build_ann_index=build_ann_index,
pg_host=os.environ.get("LLM_PG_HOST", str(section.pg_host)),

166
src/llm/evidence.py
View File

@@ -26,7 +26,7 @@ from sqlalchemy import text
from llm.config import LlmConfig
from llm.links import as_source
from pfs.codetables import is_missing_table_error
from pfs.families import detect_codes, refresh_from
from pfs.families import FAMILIES, Detection, detect_codes, refresh_from
from pfs.valuation import ValuationRow, valuation
log = logging.getLogger(__name__)
@@ -44,16 +44,91 @@ def _n(v: float | None, digits: int = 2, money: bool = False) -> str:
return f"${s}" if money else s
def cap_codes(det: Detection, max_codes: int) -> tuple[str, ...]:
"""Ruling B11 (I1's per-turn cap): *det.codes*, capped to *max_codes*
when detection turned up more than that — a wide multi-family
question can detect dozens of codes, and pricing/lineage-ing every
one blows the prompt budget on one mention. ``det.explicit`` (codes
literally named in the question) survive first; the rest fill in
family order (as ``det.families`` lists them, each family's own
codes sorted) until the cap; any code reached by neither (shouldn't
happen given how ``detect_codes`` builds ``.codes``, but kept for
safety) fills any remaining room in ``det.codes`` order. A no-op
(returns ``det.codes`` unchanged) at or under the cap — the common
case sees no behavior change. Shared by ``valuation_evidence`` and
``llm.lineage.lineage_evidence`` so a question detects the SAME
capped code list in both (and so in ``code_cited_sources``, which
receives their union)."""
if len(det.codes) <= max_codes:
return det.codes
codes_set = set(det.codes)
kept: list[str] = [c for c in det.explicit if c in codes_set]
seen = set(kept)
for fam_key in det.families:
fam = FAMILIES.get(fam_key)
if fam is None:
continue
for code in sorted(fam.codes):
if code in codes_set and code not in seen:
kept.append(code)
seen.add(code)
for code in det.codes:
if code not in seen:
kept.append(code)
seen.add(code)
capped = tuple(kept[:max_codes])
log.warning(
"code cap: dropped %d of %d detected codes (chat_codes_max=%d)",
len(det.codes) - len(capped),
len(det.codes),
max_codes,
)
return capped
def _cap_valuation_rows(
rows: Sequence[ValuationRow], explicit: Sequence[str], max_rows: int
) -> list[ValuationRow]:
"""*rows*, capped to *max_rows* (Ruling B11): explicit-code rows
survive first, then the newest vintage — selection only, the
surviving rows keep their original relative order (same pattern as
``llm.lineage._select_for_prompt``). A no-op under the cap."""
if len(rows) <= max_rows:
return list(rows)
explicit_set = set(explicit)
ranked = sorted(
range(len(rows)),
key=lambda i: (0 if rows[i].code in explicit_set else 1, -rows[i].year),
)
keep = set(ranked[:max_rows])
return [r for i, r in enumerate(rows) if i in keep]
@dataclass(frozen=True)
class ValuationEvidence:
codes: tuple[str, ...]
families: tuple[str, ...]
rows: tuple[ValuationRow, ...]
unpriced: tuple[str, ...]
#: Ruling B11: codes literally named in the question — rows for
#: these are never dropped by the prompt row cap before rows of a
#: merely-detected family-mate.
explicit: tuple[str, ...] = ()
#: ``cfg.valuation_rows_max`` at construction time — ``prompt_block``
#: is self-contained, like ``LineageEvidence.max_prompt_rows``.
max_prompt_rows: int = 24
#: ``det.wide`` (Ruling B2 family keys too big to expand into
#: ``codes``) — carried through so ``rag.stream_answer`` doesn't
#: need a third ``detect_codes`` call just to find them.
wide: tuple[str, ...] = ()
def prompt_block(self) -> str:
def prompt_block(self, max_rows: int | None = None) -> str:
"""*max_rows* overrides ``self.max_prompt_rows`` for one call —
``build_messages``'s budget trimming (Ruling B11) re-renders a
smaller block on demand without needing a new ``ValuationEvidence``."""
cap = self.max_prompt_rows if max_rows is None else max_rows
lines = [_HEADER]
for r in self.rows:
for r in _cap_valuation_rows(self.rows, self.explicit, cap):
# The status note says why an unpaid code shows no dollars, and
# the CF note says which of the year's CFs the money came from.
status = f"{r.status or '?'}" + (
@@ -173,6 +248,7 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non
det = detect_codes(question)
if not det.codes:
return None
codes = cap_codes(det, cfg.chat_codes_max)
path = _replica_path(cfg)
try:
# A chat turn runs in a threadpool worker, and DuckDB gives each
@@ -184,15 +260,23 @@ def valuation_evidence(question: str, cfg: LlmConfig) -> ValuationEvidence | Non
return None
# A whole family is six codes; showing four vintages of each would
# crowd the excerpts out of the prompt.
years = cfg.valuation_years if len(det.codes) <= 3 else 2
years = cfg.valuation_years if len(codes) <= 3 else 2
try:
rows, unpriced = valuation(cur, list(det.codes), years=years)
rows, unpriced = valuation(cur, list(codes), years=years)
except Exception as e: # noqa: BLE001
log.warning("valuation evidence skipped (query): %s", e)
return None
finally:
cur.close() # the cached parent connection stays open
return ValuationEvidence(det.codes, det.families, tuple(rows), tuple(unpriced))
return ValuationEvidence(
codes,
det.families,
tuple(rows),
tuple(unpriced),
explicit=det.explicit,
max_prompt_rows=cfg.valuation_rows_max,
wide=det.wide,
)
_CITED_SQL = text(
@@ -509,12 +593,18 @@ def _guideline_row(con: Any, path_key: str) -> tuple[int, str, str, str, str] |
return None
#: Ruling B11: manual sources never crowd more than this many slots out
#: of the prompt total, regardless of how many families were detected.
_MANUAL_SOURCES_MAX = 2
def manual_sources(
con: Any,
families: Sequence[str],
*,
per_family: int = 1,
snippet_chars: int = 900,
max_total: int = _MANUAL_SOURCES_MAX,
) -> list[dict]:
"""The AMA CPT manual's own guideline text for each detected
*family*'s CPT heading(s) — built straight from the DuckDB replica's
@@ -526,19 +616,24 @@ def manual_sources(
CPT heading ``path_key``), the newest edition with a non-empty
``guideline`` for that heading (one-level parent fallback,
``_guideline_row``) becomes one source via ``as_source``, labelled
``"CPT <year> — <heading title>"``; the source's ``date``/``url``
``"<heading title> — CPT <year>"``; the source's ``date``/``url``
come from the bib item at the edition's ``item_key`` when it
resolves, else ``f"{year}-01-01"``/``""``. At most *per_family*
sources per family (newest-edition notes are not otherwise
ordered — the first *per_family* distinct notes, alphabetically,
are used). A family with no note, or whose heading(s) have no
guideline text anywhere, contributes nothing. *con* is a cursor on
the cached replica connection (the same per-call-cursor pattern as
``valuation_evidence``). Never raises — a missing
``pfs.code_family``/``pfs.cpt_section`` table degrades to ``[]``,
same as every other replica read in this module."""
are used), and at most *max_total* across every family (Ruling B11
— a multi-family history question must not let manual sources
crowd out retrieved/cited excerpts). A family with no note, or
whose heading(s) have no guideline text anywhere, contributes
nothing. *con* is a cursor on the cached replica connection (the
same per-call-cursor pattern as ``valuation_evidence``). Never
raises — a missing ``pfs.code_family``/``pfs.cpt_section`` table
degrades to ``[]``, same as every other replica read in this
module."""
out: list[dict] = []
for family in families:
if len(out) >= max_total:
break
try:
notes = _family_notes(con, family)
except Exception as e: # noqa: BLE001
@@ -547,7 +642,7 @@ def manual_sources(
continue
kept = 0
for note in notes:
if kept >= per_family:
if kept >= per_family or len(out) >= max_total:
break
try:
row = _guideline_row(con, note)
@@ -572,7 +667,7 @@ def manual_sources(
"item_key": item_key,
"seq": str(sec_id),
"section": title,
"title": f"CPT {year} — {title}",
"title": f"{title} — CPT {year}",
"date": date,
"url": url,
}
@@ -581,9 +676,46 @@ def manual_sources(
return out
def _rule_key(s: dict) -> tuple[str, str] | None:
"""*s*'s ``(item_key, p_id)`` when it's a rule-kind source with both
— ``None`` otherwise (a rule row missing either, or any non-rule
kind), so the caller falls back to deduping on ``label`` instead."""
if s.get("kind") == "rule" and s.get("item_key") and s.get("p_id"):
return (s["item_key"], s["p_id"])
return None
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]
"""*retrieved* + whatever of *extra* isn't already present.
Ruling B13 (I2): a rule-kind source with both ``item_key`` and
``p_id`` is deduped on that pair — two rule chunks for the same FR
paragraph can carry different labels now that labels include the
volume/page (``rule_label``), and label-only dedupe would let the
same paragraph through twice under two labels. Everything else
(comments, corpus, and a rule row missing item_key/p_id) still
dedupes on ``label``, as before."""
seen_keys: set[tuple[str, str]] = set()
seen_labels: set[str] = set()
for s in retrieved:
key = _rule_key(s)
if key is not None:
seen_keys.add(key)
else:
seen_labels.add(s["label"])
out = list(retrieved)
for s in extra:
key = _rule_key(s)
if key is not None:
if key in seen_keys:
continue
seen_keys.add(key)
else:
if s["label"] in seen_labels:
continue
seen_labels.add(s["label"])
out.append(s)
return out
# Re-exported so rag.py imports both evidence functions from one place,

250
src/llm/lineage.py
View File

@@ -29,6 +29,7 @@ from __future__ import annotations
import dataclasses
import logging
import os
import threading
from dataclasses import asdict, dataclass
from typing import Any, Sequence
@@ -157,10 +158,27 @@ def _event_line(e: LineageEvent) -> str:
return line
#: Ruling B11: a priority-kind event (created/replaces/…) is only
#: guaranteed a prompt slot up to this many per code — a code with a
#: long, eventful history (many replaces/replaced_by rows) must not by
#: itself consume the whole prompt row budget.
_PRIORITY_PER_CODE_CAP = 4
def _codes_str(codes: Sequence[str]) -> str:
"""*codes*, comma-joined — or, past 8, ``"N codes (a, b, c, …)"``
with the first three (Ruling B11) — a code list this long is a
wide-family element diff, not something worth spelling out fully in
the prompt."""
if len(codes) > 8:
return f"{len(codes)} codes ({', '.join(codes[:3])}, …)"
return ", ".join(codes)
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)}"
f"in {_codes_str(d.in_codes)}; not in {_codes_str(d.not_in_codes)}"
)
@@ -171,19 +189,32 @@ def _guidance_line(g: GuidanceRef) -> str:
def _select_for_prompt(
events: Sequence[LineageEvent], max_rows: int
) -> list[LineageEvent]:
"""Every priority-kind event, plus the earliest-sorted non-priority
"""Every priority-kind event (capped at ``_PRIORITY_PER_CODE_CAP``
per code — Ruling B11), plus the earliest-sorted non-priority
events, up to *max_rows* total — priority rows are never dropped
even past the cap. Selection is by identity (``id()``), not value
equality, so it stays correct before FR urls are resolved (several
collapsed rows can otherwise compare equal on every other field).
Shared by ``LineageEvidence.prompt_block`` (final render) and
by the non-priority budget, but the whole selection is still hard-
capped at ``2 * max_rows`` (Ruling B11: a question with many
eventful codes could otherwise make priority rows alone blow the
prompt). Selection is by identity (``id()``), not value equality,
so it stays correct before FR urls are resolved (several collapsed
rows can otherwise compare equal on every other field), and render
order follows *events*' own (chronological) order. Shared by
``LineageEvidence.prompt_block`` (final render) and
``lineage_evidence`` (deciding which rows are worth an FR url
lookup — ruling: resolve links only for what the model will see)."""
priority = [e for e in events if e.kind in _PRIORITY_KINDS]
per_code: dict[str, int] = {}
priority_ids: set[int] = set()
for e in events:
if e.kind in _PRIORITY_KINDS:
n = per_code.get(e.code, 0)
if n < _PRIORITY_PER_CODE_CAP:
priority_ids.add(id(e))
per_code[e.code] = n + 1
rest = [e for e in events if e.kind not in _PRIORITY_KINDS]
budget = max(max_rows - len(priority), 0)
keep = {id(e) for e in priority} | {id(e) for e in rest[:budget]}
return [e for e in events if id(e) in keep]
budget = max(max_rows - len(priority_ids), 0)
keep = priority_ids | {id(e) for e in rest[:budget]}
selected = [e for e in events if id(e) in keep]
return selected[: 2 * max_rows]
@dataclass(frozen=True)
@@ -209,13 +240,21 @@ class LineageEvidence:
#: codes, so "what replaced G2058" surfaces G2058's own paragraphs
#: before its family-mates'.
explicit: tuple[str, ...] = ()
#: ``det.wide`` (Ruling B2 family keys too big to expand into
#: ``codes``) — carried through so ``rag.stream_answer`` doesn't
#: need a third ``detect_codes`` call just to find them.
wide: tuple[str, ...] = ()
def _prompt_events(self) -> list[LineageEvent]:
return _select_for_prompt(self.events, self.max_prompt_rows)
def _prompt_events(self, max_rows: int | None = None) -> list[LineageEvent]:
cap = self.max_prompt_rows if max_rows is None else max_rows
return _select_for_prompt(self.events, cap)
def prompt_block(self) -> str:
def prompt_block(self, max_rows: int | None = None) -> str:
"""*max_rows* overrides ``self.max_prompt_rows`` for one call —
``build_messages``'s budget trimming (Ruling B11) re-renders a
smaller block on demand without needing a new ``LineageEvidence``."""
lines = [_HEADER]
lines.extend(_event_line(e) for e in self._prompt_events())
lines.extend(_event_line(e) for e in self._prompt_events(max_rows))
if self.element_diffs:
lines.append("Element differences:")
lines.extend(_diff_line(d) for d in self.element_diffs)
@@ -240,29 +279,42 @@ class LineageEvidence:
return out
def rule_label(title: str, date_published: str, p_id: int) -> str:
"""``"CY2021 PFS final ¶1578"`` / ``"CY2027 PFS proposed ¶394"`` —
def rule_label(
title: str, date_published: str, p_id: int, *, vol: str = "", page: int = 0
) -> str:
"""``"CY2021 PFS final 85 FR 84639 ¶1578"`` when *vol*/*page* are
known (Ruling B13 — ``I2``), else the shorter
``"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"
if vol and page:
return f"CY{year} PFS {kind} {vol} FR {page}{p_id}"
return f"CY{year} PFS {kind}{p_id}"
# ── The bib Store, opened lazily once per process (sqlite, cheap) ──────
# ── The bib Store, opened lazily once PER THREAD (sqlite, cheap) ───────
#
# Ruling B14 / C1: bib.Store's sqlite connection is check_same_thread=True
# (bib.store.Store._con), and /chat runs each turn in Starlette's
# threadpool — a single process-wide Store used to be silently inert (or
# raise sqlite3.ProgrammingError) from every thread but the one that
# first opened it. threading.local() gives every worker thread its own
# Store, opened lazily on that thread's first use — no lock needed,
# threading.local() is itself the per-thread isolation.
_STORE: Any | None = None
_STORE_LOCK = threading.Lock()
_STORE_LOCAL = threading.local()
def _store() -> Any:
global _STORE
with _STORE_LOCK:
if _STORE is None:
from conf.connect import bib
store = getattr(_STORE_LOCAL, "store", None)
if store is None:
from conf.connect import bib
_STORE = bib()
return _STORE
store = bib()
_STORE_LOCAL.store = store
return store
# ── Process-global caches, all reset together whenever the replica is
@@ -297,19 +349,31 @@ class _ReplicaCache:
self._mtime = None
_ITEM_CACHE = _ReplicaCache() # item_key -> resolved bib Item (or None)
_URL_CACHE = _ReplicaCache() # (item_key, p_id) -> FR jump-link url
_ON_DEMAND_CACHE = _ReplicaCache() # code -> on-demand EventRow list
_PARAGRAPH_CACHE = _ReplicaCache() # (item_key, p_id) -> fr_anchors+items row (or None)
#: item_key -> resolved bib Item (or None); (item_key, p_id) -> FR jump-link
#: url; (item_key, p_id) -> fr_anchors+items+fr_anchor_docs row (or None).
#: Keyed on the BIB file's mtime (``_bib_mtime``), not the DuckDB
#: replica's — the two republish independently, and a lineage label/url
#: that never picks up a bib fix (a re-grabbed anchor, a corrected
#: title) is a real failure mode a long-running chat process can hit.
_ITEM_CACHE = _ReplicaCache()
_URL_CACHE = _ReplicaCache()
_PARAGRAPH_CACHE = _ReplicaCache()
#: code -> on-demand EventRow list — keyed on the DuckDB REPLICA's mtime
#: (passed explicitly by callers): on-demand lineage is derived from
#: ``pfs.rvu``/``pfs.code_event``, which live on the replica, not the
#: bib store.
_ON_DEMAND_CACHE = _ReplicaCache()
#: ``lineage_sources``'s paragraph fetches key the ``_ReplicaCache`` on a
#: constant "generation" rather than the duckdb replica's mtime: the bib
#: sqlite store (unlike the replica) is opened once per process
#: (``_store``, above) with no republish signal this module can key on,
#: so a paragraph, once fetched, is cached for the life of the process —
#: exactly the pattern the item/url caches use, minus the invalidation
#: trigger they don't have here either.
_PARAGRAPH_CACHE_GEN = 0
def _bib_mtime(store: Any) -> int:
"""The bib sqlite file's mtime, for ``_ITEM_CACHE``/``_URL_CACHE``/
``_PARAGRAPH_CACHE`` — mirrors ``llm.rag._docket_year``'s own stat of
the same file. Never raises (a store with no ``_db_path``, or a
broken one, just never gets a cache hit)."""
try:
return os.stat(store._db_path).st_mtime_ns
except Exception:
return -1
def _cached_item(store: Any, item_key: str, mtime: int) -> Any:
@@ -330,11 +394,25 @@ def _source_label(store: Any, item_key: str, p_id: int, mtime: int) -> str:
bare item key when the item can't be fetched (unresolvable/missing
item must never break the chat over a label). The resolved
``Store.get`` item is cached per ``item_key`` — the same handful of
rules back dozens of events/diffs/guidance rows in one turn."""
rules back dozens of events/diffs/guidance rows in one turn.
Ruling B13 (I2): when *p_id* resolves to a real FR paragraph
(``_lineage_paragraph``, itself cached per ``(item_key, p_id)``),
its volume/page are passed to ``rule_label`` so the label is unique
across paragraphs that would otherwise share one — a p_id alone
collides across rules 42% of the time. Falls back to the shorter
form when the paragraph can't be found (guidance citing a CPT
edition item with ``p_id=0``, an unanchored row, …)."""
item = _cached_item(store, item_key, mtime)
if item is None:
return item_key or "?"
return rule_label(item.title, item.date_published, p_id)
vol, page = "", 0
if p_id:
para = _lineage_paragraph(store, item_key, p_id)
if para is not None:
vol = str(para.get("fr_volume") or "")
page = para.get("page") or 0
return rule_label(item.title, item.date_published, p_id, vol=vol, page=page)
def _fr_url(store: Any, item_key: str, p_id: int, mtime: int) -> str:
@@ -401,15 +479,41 @@ def _resolve_prompt_urls(
return tuple(out)
def _pref_key(r: EventRow) -> tuple[int, int]:
"""Collapse preference within a group: an anchored FR row first,
then the lowest ``p_id`` — shared by both ``_collapse`` passes."""
return (0 if (r.source == "fr" and r.anchored) else 1, r.p_id)
def _is_proposed(store: Any, item_key: str, mtime: int) -> bool:
"""True when *item_key*'s title contains "proposed" — ``False`` (not
"unknown") for an empty *item_key* or a *store*-less caller, which
keeps every ``_pref_key`` caller that doesn't have a store (existing
tests, mostly) exactly as deterministic as before Ruling B12."""
if store is None or not item_key:
return False
item = _cached_item(store, item_key, mtime)
if item is None:
return False
return "proposed" in (item.title or "").lower()
def _collapse(rows: Sequence[EventRow]) -> list[EventRow]:
def _pref_key(
r: EventRow, store: Any = None, mtime: int = 0
) -> tuple[bool, bool, int, str]:
"""Collapse preference within a group (Ruling B12): an anchored FR
row first, then a final rule over a proposed one, then the lowest
``p_id``, then ``item_key`` (a total order — the last two fields
exist only to make a fully-tied pair deterministic). Shared by both
``_collapse`` passes. *store*/*mtime* are optional so callers with
no store (every pre-B12 test) get the same ordering as before on
anything that isn't tied on ``(anchored_fr, p_id)`` alone."""
anchored_fr = r.source == "fr" and r.anchored
proposed = _is_proposed(store, r.item_key, mtime) if r.source == "fr" else False
return (not anchored_fr, proposed, r.p_id, r.item_key)
def _collapse(
rows: Sequence[EventRow], store: Any = None, mtime: int = 0
) -> list[EventRow]:
"""One representative row per ``(code, year, kind, from_codes,
to_codes)`` — prefer an anchored FR row, then the lowest ``p_id``.
to_codes)`` — prefer an anchored FR row, then a final rule over a
proposed one, then the lowest ``p_id``, then ``item_key`` (Ruling
B12 — ``_pref_key``).
Ruling B5: "point in time" kinds (see ``_EARLIEST_ONLY_KINDS``)
collapse further to a single row per ``(code, kind)`` at the
@@ -417,13 +521,20 @@ def _collapse(rows: Sequence[EventRow]) -> list[EventRow]:
year's rows when more than one survives. Directional/repeatable
kinds keep one row per year.
Sorted by ``(year, code, kind order)``.
Sorted by ``(year, code, kind order)``. *store*/*mtime* resolve
each row's rule title for the proposed/final tiebreak (cached per
``item_key`` — ``_is_proposed``); omitted, every row is treated as
not-proposed, same as before Ruling B12.
"""
def key(r: EventRow) -> tuple[bool, bool, int, str]:
return _pref_key(r, store, mtime)
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=_pref_key) for grp in groups.values()]
k = (r.code, r.year, r.kind, r.from_codes, r.to_codes)
groups.setdefault(k, []).append(r)
reps = [min(grp, key=key) for grp in groups.values()]
singleton: dict[tuple[str, str], list[EventRow]] = {}
other: list[EventRow] = []
@@ -433,7 +544,7 @@ def _collapse(rows: Sequence[EventRow]) -> list[EventRow]:
else:
other.append(r)
earliest = [
min((r for r in grp if r.year == min(x.year for x in grp)), key=_pref_key)
min((r for r in grp if r.year == min(x.year for x in grp)), key=key)
for grp in singleton.values()
]
@@ -597,9 +708,13 @@ def _lineage_paragraph(store: Any, item_key: str, p_id: int) -> dict | None:
``fr_anchor_docs`` (``html_url``/``fr_volume`` — the fields
``links._rule`` and ``source._build_rule_doc`` use to build a rule
chunk's own metadata) for one ``(item_key, p_id)``. ``None`` when
the paragraph doesn't exist. Cached per key for the life of the
process (``_PARAGRAPH_CACHE``, see its comment above)."""
cached = _PARAGRAPH_CACHE.get((item_key, p_id), _PARAGRAPH_CACHE_GEN)
the paragraph doesn't exist. Cached per key, invalidated when the
bib file's mtime changes (``_bib_mtime`` — see ``_ITEM_CACHE``'s
comment above; this function has no mtime parameter of its own
since some callers, e.g. ``lineage_sources`` from ``rag.py``, have
no mtime to hand it)."""
mtime = _bib_mtime(store)
cached = _PARAGRAPH_CACHE.get((item_key, p_id), mtime)
if cached is not _MISSING:
return cached
row = (
@@ -616,7 +731,7 @@ def _lineage_paragraph(store: Any, item_key: str, p_id: int) -> dict | None:
.fetchone()
)
result = dict(row) if row is not None else None
_PARAGRAPH_CACHE.put((item_key, p_id), _PARAGRAPH_CACHE_GEN, result)
_PARAGRAPH_CACHE.put((item_key, p_id), mtime, result)
return result
@@ -723,6 +838,11 @@ def lineage_sources(
return []
out: list[dict] = []
seen: set[tuple[str, int]] = set()
# Ruling B11: _LINEAGE_SOURCES_HARD_CAP is an absolute ceiling
# regardless of a config-raised max_items — the event loop must
# enforce it itself, not rely on the element-diff loop below (which
# only ever runs after this one).
event_cap = min(max_items, _LINEAGE_SOURCES_HARD_CAP)
for e in candidates:
key = (e.item_key, e.p_id)
if key in seen:
@@ -735,7 +855,7 @@ def lineage_sources(
continue
if src is not None:
out.append(src)
if len(out) >= max_items:
if len(out) >= event_cap:
break
added = 0
@@ -810,6 +930,7 @@ def lineage_evidence(question: str, cfg: LlmConfig) -> LineageEvidence | None:
det = detect_codes(question)
if not det.codes:
return None
codes = evidence.cap_codes(det, cfg.chat_codes_max)
path = evidence._replica_path(cfg)
try:
con = evidence._connect(path)
@@ -819,23 +940,27 @@ def lineage_evidence(question: str, cfg: LlmConfig) -> LineageEvidence | None:
cur = con.cursor()
try:
store = _store()
mtime = evidence._mtime(path)
mtime = evidence._mtime(path) # replica mtime — _ON_DEMAND_CACHE only
bib_mtime = _bib_mtime(store) # bib mtime — item/url/paragraph caches
raw_events = _collect_events(
cur, store, mtime, det.codes, cfg.lineage_on_demand_max
cur, store, mtime, codes, cfg.lineage_on_demand_max
)
events_list = [
_to_lineage_event(store, r, mtime) for r in _collapse(raw_events)
_to_lineage_event(store, r, bib_mtime)
for r in _collapse(raw_events, store, bib_mtime)
]
events = _resolve_prompt_urls(store, events_list, mtime, cfg.lineage_max_rows)
element_diffs, elements_note = _element_diffs(cur, det.codes, store, mtime)
guidance = _collect_guidance(cur, store, det.families, det.wide, mtime)
events = _resolve_prompt_urls(
store, events_list, bib_mtime, cfg.lineage_max_rows
)
element_diffs, elements_note = _element_diffs(cur, codes, store, bib_mtime)
guidance = _collect_guidance(cur, store, det.families, det.wide, bib_mtime)
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,
codes=codes,
families=det.families,
events=events,
element_diffs=element_diffs,
@@ -843,4 +968,5 @@ def lineage_evidence(question: str, cfg: LlmConfig) -> LineageEvidence | None:
elements_note=elements_note,
max_prompt_rows=cfg.lineage_max_rows,
explicit=det.explicit,
wide=det.wide,
)

298
src/llm/rag.py
View File

@@ -51,7 +51,6 @@ from llm.links import as_source
from llm.pool import HostPool, PoolEmbeddings, pick_model
from llm.rerank import Hit, _similarity, blend, filter_since
from pfs.descriptors import rule_year_of
from pfs.families import detect_codes
log = logging.getLogger(__name__)
@@ -113,27 +112,38 @@ def _docket_rule_year(comment_end_date: str) -> int:
def _docket_year(docket_id: str) -> int:
"""The PFS rule year *docket_id* belongs to, 0 when unknown."""
"""The PFS rule year *docket_id* belongs to, 0 when unknown.
Uses the same thread-local ``_bib_store()`` accessor as ``llm.lineage``
(C1/Ruling B14) — each thread opens its own bib.Store. Never raises:
a failure opening the store, reading its mtime, or querying its
dockets is logged and degrades to 0, which ``era_of`` already falls
back on to the hit's own date year — this must never escape and
break era-balanced retrieval over one bad docket lookup."""
global _docket_era_mtime
if not docket_id:
return 0
store = _bib_store()
try:
mtime: object = os.stat(store._db_path).st_mtime_ns
except Exception:
mtime = None
with _docket_era_lock:
if _docket_era_mtime != mtime:
_docket_era_map.clear()
_docket_era_map.update(
{
d.id: _docket_rule_year(d.comment_end_date)
for d in store.dockets()
if d.comment_end_date
}
)
_docket_era_mtime = mtime
return _docket_era_map.get(docket_id, 0)
store = _bib_store()
try:
mtime: object = os.stat(store._db_path).st_mtime_ns
except Exception:
mtime = None
with _docket_era_lock:
if _docket_era_mtime != mtime:
_docket_era_map.clear()
_docket_era_map.update(
{
d.id: _docket_rule_year(d.comment_end_date)
for d in store.dockets()
if d.comment_end_date
}
)
_docket_era_mtime = mtime
return _docket_era_map.get(docket_id, 0)
except Exception as e: # noqa: BLE001 — a bad docket lookup must not break era_of
log.warning("docket era lookup failed (%s): %s", docket_id, e)
return 0
def _year_of(iso_date: str) -> int:
@@ -183,6 +193,17 @@ def era_balance(hits: list[Hit], *, per_era: int, top_n: int) -> list[Hit]:
for h in ranked:
by_era.setdefault(era_of(h), []).append(h)
eras_desc = sorted(by_era, reverse=True)
if len(eras_desc) > top_n > 1:
# I3: eras_desc has more distinct eras than top_n — the
# round-robin below visits it newest-first every round, so
# without this the first round alone (>= top_n eras) fills the
# output and the oldest eras never survive the final top_n cut.
# Pick top_n eras evenly spaced across the (already
# newest-first) list, by index, so both ends — the oldest and
# the newest era — always survive.
n = len(eras_desc)
idxs = sorted({round(i * (n - 1) / (top_n - 1)) for i in range(top_n)})
eras_desc = [eras_desc[i] for i in idxs]
capped = {era: by_era[era][:per_era] for era in eras_desc}
selected: list[Hit] = []
@@ -210,29 +231,30 @@ def _resolve_mode(question: str, mode: str) -> str:
_SYSTEM = (
"You answer questions about CMS rulemaking using ONLY the excerpts "
"provided below. Excerpts come from three kinds of sources: public "
"You answer questions about CMS rulemaking using ONLY the material "
"provided below — the excerpts, and the Lineage and Valuation sections "
"when present. Excerpts come from three kinds of sources: public "
"comments submitted to regulations.gov dockets, Federal Register rules "
"(proposed and final), and a reference library (journal articles, CMS "
"manuals, regulations, agency documents). Each excerpt is prefixed with "
"its citation label in square brackets and its kind and date. When you "
"use an excerpt, cite its label exactly, e.g. [CMS-2026-2377-3438] or "
"[91 FR 43949 ¶4]. Prefer the most recent comments when excerpts "
"conflict or describe a changing position, and say what year a "
"statement comes from when it matters. If the excerpts do not contain "
"the answer, say you don't have information on that in the indexed "
"library — do not invent facts."
" A Valuation section may follow the excerpts. Its rows are "
"authoritative for RVUs, conversion factors and payment amounts; when "
"you state any of those numbers, cite the row's bracketed label exactly, "
"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."
"conflict or describe a changing position, unless the question is about "
"history or a timeline — say what year a statement comes from when it "
"matters. If none of the provided material contains the answer, say "
"you don't have information on that in the indexed library — do not "
"invent facts."
" A Lineage section may also follow the excerpts. 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."
" A Valuation section may also follow. Its rows are authoritative for "
"RVUs, conversion factors and payment amounts; when you state any of "
"those numbers, cite the row's bracketed label exactly, 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."
)
@@ -305,15 +327,31 @@ def retrieve(
return [_source(h) for h in ranked]
def build_messages(
#: Ruling B11 drop-order floors — how many of each source block survive
#: once ``build_messages``'s budget trimming kicks in. Sources beyond
#: these are dropped first (lineage-source excerpts), through last
#: (cited); valuation/lineage prompt ROWS are trimmed only after every
#: source block is already at its floor.
_BUDGET_LINEAGE_SRC_KEEP = 4
_BUDGET_RETRIEVED_KEEP = 6
_BUDGET_CITED_KEEP = 8
_BUDGET_MANUAL_KEEP = 1
_BUDGET_VALUATION_ROWS_KEEP = 12
def _user_message(
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 Lineage block and a Valuation block — in that
order — between the excerpts and the question."""
retrieved: list[dict],
cited: list[dict],
lineage_src: list[dict],
manual: list[dict],
rest: list[dict],
evidence: ValuationEvidence | None,
lineage: LineageEvidence | None,
valuation_rows_cap: int | None,
lineage_rows_cap: int | None,
) -> str:
sources = retrieved + cited + lineage_src + manual + rest
if sources:
context = "\n\n".join(
f"[{s['label']}] ({s.get('kind', '')}, {s.get('date', '') or 'undated'}) "
@@ -324,14 +362,149 @@ def build_messages(
context = "(no relevant excerpts found)"
parts = [f"Excerpts:\n\n{context}"]
if lineage is not None:
parts.append(lineage.prompt_block())
parts.append(lineage.prompt_block(lineage_rows_cap))
if evidence is not None:
parts.append(evidence.prompt_block())
parts.append(evidence.prompt_block(valuation_rows_cap))
parts.append(f"Question: {question}")
user = "\n\n".join(parts)
return "\n\n".join(parts)
def build_messages(
question: str,
sources: list[dict],
evidence: ValuationEvidence | None = None,
lineage: LineageEvidence | None = None,
*,
n_cited: int = 0,
n_lineage_sources: int = 0,
n_manual: int = 0,
n_retrieved: int | None = None,
budget_chars: int | None = None,
) -> list[dict]:
"""Grounded chat messages: system rules + question with excerpts and,
when available, a Lineage block and a Valuation block — in that
order — between the excerpts and the question.
*sources* is assumed laid out contiguously in the order
``stream_answer`` appends them — retrieved, then cited
(``code_cited_sources``), then lineage-source excerpts
(``llm.lineage.lineage_sources``), then manual (``manual_sources``),
then anything else — with *n_cited*/*n_lineage_sources*/*n_manual*
giving each block's length (0 by default: every existing caller
that passes one flat, unlabelled ``sources`` list behaves exactly
as before). *n_retrieved* defaults to whatever of *sources* isn't
claimed by the other three counts.
When *budget_chars* is given (Ruling B11) and the assembled user
message exceeds it, blocks are trimmed in this order until it fits
(or every step is exhausted): lineage-source excerpts beyond 4,
retrieved beyond 6, cited beyond 8, manual beyond 1, valuation
prompt rows beyond 12, lineage prompt rows beyond
``lineage.max_prompt_rows``. Each actual drop is logged with the
block name and the over-budget char count. ``budget_chars=None``
(the default) never trims — every pre-B11 caller is unaffected.
"""
total = len(sources)
if n_retrieved is None:
n_retrieved = max(total - n_cited - n_lineage_sources - n_manual, 0)
retrieved = sources[:n_retrieved]
cited = sources[n_retrieved : n_retrieved + n_cited]
lsrc_end = n_retrieved + n_cited + n_lineage_sources
lineage_src = sources[n_retrieved + n_cited : lsrc_end]
manual = sources[lsrc_end : lsrc_end + n_manual]
rest = sources[lsrc_end + n_manual :]
valuation_cap = evidence.max_prompt_rows if evidence is not None else None
lineage_cap = lineage.max_prompt_rows if lineage is not None else None
def content() -> str:
return _user_message(
question,
retrieved,
cited,
lineage_src,
manual,
rest,
evidence,
lineage,
valuation_cap,
lineage_cap,
)
if budget_chars is not None and len(content()) > budget_chars:
log.info(
"prompt budget: %d chars over %d — trimming",
len(content()) - budget_chars,
budget_chars,
)
# Sources, in drop priority: lineage-source excerpts first
# (the model can still read the events' own labels even
# without their paragraph text), through cited last.
source_steps: list[tuple[str, list[dict], int]] = [
("lineage-source excerpts", lineage_src, _BUDGET_LINEAGE_SRC_KEEP),
("retrieved", retrieved, _BUDGET_RETRIEVED_KEEP),
("cited", cited, _BUDGET_CITED_KEEP),
("manual", manual, _BUDGET_MANUAL_KEEP),
]
for name, block, keep in source_steps:
if len(content()) <= budget_chars:
break
if len(block) <= keep:
continue
before = len(content())
dropped = len(block) - keep
block[:] = block[:keep]
if name == "lineage-source excerpts":
lineage_src = block
elif name == "retrieved":
retrieved = block
elif name == "cited":
cited = block
else:
manual = block
log.info(
"prompt budget: dropped %d %s source(s) (%d chars over %d)",
dropped,
name,
before - budget_chars,
budget_chars,
)
# Then the prompt ROWS the evidence blocks render, last resort:
# valuation to a flat 12, lineage down to (at most) its own
# configured lineage_max_rows — overriding the softer "priority
# rows may run up to 2x that" leeway _select_for_prompt allows
# by default (``llm.lineage._select_for_prompt``'s own hard cap
# is ``2 * max_rows``, so requesting half of the target keeps
# the actual rendered count at or under it).
if (
len(content()) > budget_chars
and evidence is not None
and (valuation_cap is None or valuation_cap > _BUDGET_VALUATION_ROWS_KEEP)
):
before = len(content())
valuation_cap = _BUDGET_VALUATION_ROWS_KEEP
log.info(
"prompt budget: capped valuation rows to %d (%d chars over %d)",
valuation_cap,
before - budget_chars,
budget_chars,
)
if len(content()) > budget_chars and lineage is not None:
target = lineage.max_prompt_rows
half = max(target // 2, 1)
if lineage_cap is None or lineage_cap > half:
before = len(content())
lineage_cap = half
log.info(
"prompt budget: capped lineage rows to <= %d (%d chars over %d)",
target,
before - budget_chars,
budget_chars,
)
return [
{"role": "system", "content": _SYSTEM},
{"role": "user", "content": user},
{"role": "user", "content": content()},
]
@@ -352,6 +525,9 @@ def _manual_sources(families: Sequence[str], cfg: LlmConfig) -> list[dict]:
return []
try:
return manual_sources(cur, families)
except Exception as e: # noqa: BLE001 — a manual-source bug must not break the chat
log.warning("manual sources skipped (query): %s", e)
return []
finally:
cur.close() # the cached parent connection stays open
@@ -391,6 +567,8 @@ def stream_answer(
"""
resolved_mode = _resolve_mode(question, mode)
sources = retrieve(question, cfg=cfg, pool=pool, since=since, mode=mode)
n_retrieved = len(sources)
n_cited = n_lineage_sources = n_manual = 0
lineage = lineage_evidence(question, cfg)
evidence = valuation_evidence(question, cfg)
codes: tuple[str, ...] = evidence.codes if evidence is not None else ()
@@ -408,7 +586,9 @@ def stream_answer(
max_total=cfg.code_cited_max,
by_docket=(resolved_mode == "timeline"),
)
before = len(sources)
sources = merge_sources(sources, cited)
n_cited = len(sources) - before
if lineage is not None:
# The FR paragraphs behind the events selected into the
# lineage prompt block, so the model can read what it's
@@ -421,19 +601,39 @@ def stream_answer(
_bib_store(), lineage, max_items=cfg.lineage_sources_max
)
sources, lrows = _reconcile_labels(sources, lrows)
before = len(sources)
sources = merge_sources(sources, lrows)
n_lineage_sources = len(sources) - before
# The CPT manual's own heading text for a detected family — a
# wide family (its code list too big to expand, MAX_FAMILY_EXPAND)
# is excluded, and these rows never count toward code_cited_max.
wide = set(detect_codes(question).wide)
# ``wide`` comes off whichever of evidence/lineage already ran
# detect_codes — no third detect_codes call for this question.
wide: set[str] = set()
if evidence is not None:
wide |= set(evidence.wide)
if lineage is not None:
wide |= set(lineage.wide)
narrow_families = tuple(f for f in families if f not in wide)
before = len(sources)
sources = merge_sources(sources, _manual_sources(narrow_families, cfg))
n_manual = len(sources) - before
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, lineage)
messages = build_messages(
question,
sources,
evidence,
lineage,
n_retrieved=n_retrieved,
n_cited=n_cited,
n_lineage_sources=n_lineage_sources,
n_manual=n_manual,
budget_chars=cfg.chat_num_ctx * 3,
)
with pool.acquire_generation() as host, httpx.Client(timeout=_TIMEOUT) as client:
model = pick_model(cfg, pool, host)
with client.stream(

2
stack.toml
View File

@@ -134,6 +134,8 @@ code_cited_max = 12 # cited excerpts kept after round-robin
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
lineage_sources_max = 8 # FR paragraphs fetched as sources for events selected into the lineage prompt block, outside code_cited_max (Ruling B10: + up to 2 element-diff anchors, hard cap 10 total)
chat_codes_max = 24 # Ruling B11: per-turn cap on detected codes (explicit codes first, then family order) — applied before valuation/lineage/code-cited queries
valuation_rows_max = 24 # Ruling B11: valuation rows kept in the prompt block (explicit codes first, then newest vintage); the SSE payload always carries every row
[llm.k_per_kind] # over-fetched ×3 per kind, then re-ranked
comment = 8

27
tests/conftest.py
View File

@@ -11,11 +11,35 @@ Fixtures are organized by domain:
from __future__ import annotations
import re
import shutil
import tempfile
from pathlib import Path
import polars as pl
import pytest
# ── pfs.families registry fixture (shared: pfs + llm tests) ──────────────────
# Any test that calls ``refresh_from``/mutates the live ``pfs.families.FAMILIES``
# registry (directly, or indirectly by opening a real replica — I5) must not
# leave it clobbered for the rest of the session.
@pytest.fixture
def restore_families():
"""Snapshot ``pfs.families.FAMILIES`` and restore it (and the compiled
phrase/code index built off it) in teardown, even if the test body
raises."""
from pfs import families as mod
before = dict(mod.FAMILIES)
try:
yield mod
finally:
mod.FAMILIES.clear()
mod.FAMILIES.update(before)
mod.rebuild_index()
# ── Perf hooks — auto-file Gitea issues on test failure/skip in CI ───────────
# Only active when STACK_FILE_TEST_ISSUES=true (set in CI workflows).
@@ -28,9 +52,6 @@ except ImportError:
# create_db() takes ~10s. By creating once and copying per-test, we cut
# cumulative DB creation from ~30min to ~10s across 200+ Zotero tests.
import shutil
import tempfile
@pytest.fixture(scope="session")
def _zotero_template_db():

140
tests/llm/test_evidence.py
View File

@@ -232,6 +232,97 @@ class TestValuationEvidence:
assert mock_val.call_args.kwargs == {"years": 2}
class TestCapCodes:
"""``cap_codes`` — Ruling B11 per-turn code cap."""
def test_noop_at_or_under_the_cap(self):
from pfs.families import Detection
det = Detection(codes=("A", "B"), families=(), explicit=("A",), wide=())
assert evidence.cap_codes(det, 2) is det.codes
def test_explicit_codes_survive_first(self):
from pfs.families import Detection
det = Detection(
codes=tuple(f"{i:05d}" for i in range(10)),
families=(),
explicit=("00009", "00005"),
wide=(),
)
capped = evidence.cap_codes(det, 3)
assert len(capped) == 3
assert {"00009", "00005"} <= set(capped)
def test_remaining_slots_fill_in_family_order(self, monkeypatch):
from pfs.families import Detection, Family
fam_a = Family("FAMA", "Family A", ("A1", "A2"), ())
fam_b = Family("FAMB", "Family B", ("B1", "B2"), ())
monkeypatch.setattr(evidence, "FAMILIES", {"FAMA": fam_a, "FAMB": fam_b})
det = Detection(
codes=("A1", "A2", "B1", "B2"),
families=("FAMB", "FAMA"), # detected order — B before A
explicit=(),
wide=(),
)
capped = evidence.cap_codes(det, 3)
assert capped == ("B1", "B2", "A1")
def test_logs_the_drop_count(self, caplog):
from pfs.families import Detection
det = Detection(
codes=tuple(f"{i:05d}" for i in range(5)),
families=(),
explicit=(),
wide=(),
)
evidence.cap_codes(det, 2)
assert "dropped 3 of 5" in caplog.text
class TestValuationRowCap:
"""Ruling B11: ``ValuationEvidence.prompt_block`` renders only the
capped rows (explicit codes first, then newest vintage); ``payload``
always carries every row."""
def _rows(self):
return [
_row(code="G0556", year=2023, label="[Y2023]"),
_row(code="G0556", year=2024, label="[Y2024]"),
_row(code="G0557", year=2025, label="[Y2025]"),
_row(code="G0558", year=2026, label="[Y2026]"),
]
def test_payload_keeps_every_row_regardless_of_cap(self):
ev = ValuationEvidence(
("G0556", "G0557", "G0558"), (), tuple(self._rows()), (), max_prompt_rows=2
)
assert len(ev.payload()["rows"]) == 4
def test_prompt_block_caps_and_prefers_explicit_then_newest_vintage(self):
ev = ValuationEvidence(
("G0556", "G0557", "G0558"),
(),
tuple(self._rows()),
(),
explicit=("G0557",),
max_prompt_rows=2,
)
block = ev.prompt_block()
assert "[Y2025]" in block # explicit code's row survives
assert "[Y2026]" in block # newest vintage among the rest
assert "[Y2023]" not in block and "[Y2024]" not in block
def test_max_rows_override_for_one_call(self):
ev = ValuationEvidence(
("G0556", "G0557", "G0558"), (), tuple(self._rows()), (), max_prompt_rows=24
)
assert ev.prompt_block().count("\n") + 1 == 5 # header + 4 rows, no cap
assert ev.prompt_block(max_rows=1).count("[Y") == 1
class TestWarm:
"""#699 ruling B3: warm the replica connection ahead of the first
real request — at API startup and again (a no-op once cached) at the
@@ -871,6 +962,43 @@ class TestMergeSources:
assert [s["label"] for s in merge_sources(a, b)] == ["X", "Y", "Z"]
assert merge_sources(a, b)[1]["score"] == 0.8
def test_rule_rows_dedupe_by_item_key_p_id_even_with_different_labels(self):
# Ruling B13 (I2): two rule chunks for the same paragraph can now
# carry different labels (a retrieved "85 FR 84639 ¶12" vs a
# lineage "CY2021 PFS final 85 FR 84639 ¶1578") — the pair must
# still dedupe on (item_key, p_id), not slip through on label.
a = [
{
"label": "85 FR 84639 ¶12",
"kind": "rule",
"item_key": "YBM4IZUS",
"p_id": "1578",
}
]
b = [
{
"label": "CY2021 PFS final 85 FR 84639 ¶1578",
"kind": "rule",
"item_key": "YBM4IZUS",
"p_id": "1578",
}
]
out = merge_sources(a, b)
assert len(out) == 1
assert out[0]["label"] == "85 FR 84639 ¶12"
def test_rule_row_missing_item_key_or_p_id_falls_back_to_label(self):
a = [{"label": "some rule", "kind": "rule", "item_key": "", "p_id": ""}]
b = [{"label": "some rule", "kind": "rule", "item_key": "K", "p_id": ""}]
c = [{"label": "other", "kind": "rule", "item_key": "K", "p_id": ""}]
out = merge_sources(merge_sources(a, b), c)
assert [s["label"] for s in out] == ["some rule", "other"]
def test_non_rule_kinds_still_dedupe_by_label(self):
a = [{"label": "L", "kind": "comment", "item_key": "K", "p_id": ""}]
b = [{"label": "L", "kind": "comment", "item_key": "K", "p_id": ""}]
assert len(merge_sources(a, b)) == 1
def _cpt_section_row(con, year, item_key, sec_id, title, path_key, guideline):
con.execute(
@@ -978,7 +1106,7 @@ class TestManualSources:
assert len(out) == 1
s = out[0]
assert s["kind"] == "corpus"
assert s["title"] == "CPT 2024 — Chronic Care Management Services"
assert s["title"] == "Chronic Care Management Services — CPT 2024"
assert s["item_key"] == "GQGTPGYV" and s["seq"] == "S1"
assert s["section"] == "Chronic Care Management Services"
assert s["date"] == "2023-11-16"
@@ -989,7 +1117,7 @@ class TestManualSources:
out = manual_sources(cpt_replica, ["PARENTONLY"])
assert len(out) == 1
s = out[0]
assert s["title"] == "CPT 2023 — General"
assert s["title"] == "General — CPT 2023"
assert s["item_key"] == "PARED0001" and s["seq"] == "S3"
assert "General guidance text" in s["snippet"]
# unresolved bib item (not in the fake store) — falls back to
@@ -1023,6 +1151,14 @@ class TestManualSources:
assert len(manual_sources(cpt_replica, ["CCM"], per_family=1)) == 1
assert len(manual_sources(cpt_replica, ["CCM"], per_family=2)) == 2
def test_total_cap_across_families(self, cpt_replica):
# Ruling B11: even with three families each good for a source,
# manual_sources never returns more than max_total (default 2).
_family_row(cpt_replica, "PARENTONLY2", PARENT_PATH, code="99498")
out = manual_sources(cpt_replica, ["CCM", "PARENTONLY", "PARENTONLY2"])
assert len(out) == 2
assert len(manual_sources(cpt_replica, ["CCM", "PARENTONLY"], max_total=1)) == 1
def test_snippet_chars_bounds_the_text_passed_to_as_source(self, cpt_replica):
out = manual_sources(cpt_replica, ["CCM"], snippet_chars=10)
assert len(out[0]["snippet"]) <= 10

329
tests/llm/test_lineage.py
View File

@@ -5,7 +5,9 @@ from __future__ import annotations
import json
import sqlite3
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import patch
@@ -61,7 +63,7 @@ def _clear_caches():
def _reset():
evidence._REPLICA = None
lineage._STORE = None
lineage._STORE_LOCAL = threading.local()
lineage._ITEM_CACHE.reset()
lineage._URL_CACHE.reset()
lineage._ON_DEMAND_CACHE.reset()
@@ -232,6 +234,29 @@ class TestRuleLabel:
== "CY2027 PFS proposed ¶394"
)
def test_final_with_volume_and_page(self):
assert (
rule_label(
"Medicare Program; CY 2021 Payment Policies Under the PFS",
"2020-12-28",
1578,
vol="85",
page=84639,
)
== "CY2021 PFS final 85 FR 84639 ¶1578"
)
def test_missing_page_falls_back_to_the_short_form(self):
assert (
rule_label(
"Medicare Program; CY 2021 Payment Policies Under the PFS",
"2020-12-28",
1578,
vol="85",
)
== "CY2021 PFS final ¶1578"
)
class TestCollapse:
def test_prefers_anchored_fr_over_rvu_for_the_same_group(self):
@@ -354,6 +379,73 @@ class TestCollapse:
(2022, "G2214"),
]
def test_final_beats_proposed_on_a_tie(self, store):
"""Ruling B12: two rows tied on (code, year, kind, from, to,
p_id) — one from a proposed rule (ZPROP2027), one from a final
rule (YBM4IZUS) — the final rule wins."""
rows = [
_ev(
"99439",
2021,
"created",
item_key="ZPROP2027",
p_id=999,
source="fr",
anchored=True,
),
_ev(
"99439",
2021,
"created",
item_key="YBM4IZUS",
p_id=999,
source="fr",
anchored=True,
),
]
collapsed = lineage._collapse(rows, store, 0)
assert len(collapsed) == 1
assert collapsed[0].item_key == "YBM4IZUS"
def test_fully_tied_pair_decided_by_item_key(self, store):
"""Neither row resolves to a proposed title (both unknown item
keys) — the tie is broken deterministically by item_key."""
rows = [
_ev(
"99439",
2021,
"created",
item_key="ZZZZ0000",
p_id=999,
source="fr",
anchored=True,
),
_ev(
"99439",
2021,
"created",
item_key="AAAA0000",
p_id=999,
source="fr",
anchored=True,
),
]
collapsed = lineage._collapse(rows, store, 0)
assert len(collapsed) == 1
assert collapsed[0].item_key == "AAAA0000"
def test_no_store_never_prefers_proposed_but_stays_deterministic(self):
"""Without a store (the pre-B12 call shape), every row is
treated as not-proposed — the anchored/p_id/item_key ordering
still applies and never raises."""
rows = [
_ev("99439", 2021, "created", item_key="B", p_id=999, source="fr"),
_ev("99439", 2021, "created", item_key="A", p_id=999, source="fr"),
]
collapsed = lineage._collapse(rows)
assert len(collapsed) == 1
assert collapsed[0].item_key == "A"
class TestElementDiffs:
def test_no_diffs_when_fewer_than_two_codes_have_elements(self, con):
@@ -384,7 +476,9 @@ class TestElementDiffs:
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
assert (
consent.label == "CY2021 PFS final 85 FR 84639 ¶1578"
) # from the anchored 99490 row
telehealth = by_pair[("telehealth", "yes")]
assert telehealth.in_codes == ("G0556",)
assert telehealth.not_in_codes == ("99490", "99491")
@@ -478,7 +572,22 @@ class TestGuidance:
)
out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
assert out[0].url == ""
assert out[0].label == "CY2021 PFS final ¶686"
assert out[0].label == "CY2021 PFS final 85 FR 84547 ¶686"
class TestCodesStr:
"""Ruling B11: element-diff code lists compact past 8 codes."""
def test_short_list_joined_plainly(self):
assert lineage._codes_str(("A", "B", "C")) == "A, B, C"
def test_exactly_eight_stays_plain(self):
codes = tuple(f"C{i}" for i in range(8))
assert lineage._codes_str(codes) == ", ".join(codes)
def test_over_eight_compacts_to_count_and_first_three(self):
codes = tuple(f"C{i}" for i in range(9))
assert lineage._codes_str(codes) == "9 codes (C0, C1, C2, …)"
class TestPromptBlock:
@@ -618,6 +727,72 @@ class TestPromptBlock:
for kind in ("created", "replaces", "replaced_by", "deleted"):
assert kind in rendered
def test_priority_rows_capped_at_four_per_code(self):
"""Ruling B11: a single code with 6 priority-kind rows only ever
gets 4 of them into the prompt."""
priorities = [
LineageEvent(
"99490",
2020 + i,
"replaces",
(),
(),
f"[L{i}]",
"K",
i,
0,
"",
"fr",
True,
"",
)
for i in range(6)
]
ev = LineageEvidence(
codes=("99490",),
families=(),
events=tuple(priorities),
element_diffs=(),
guidance=(),
max_prompt_rows=25,
)
rendered = ev.prompt_block()
assert sum(rendered.count(f"[L{i}]") for i in range(6)) == 4
def test_hard_cap_at_twice_max_rows(self):
"""Ruling B11: many codes each with priority rows can push the
priority selection itself past max_rows — the absolute ceiling
is 2 * max_rows."""
priorities = [
LineageEvent(
f"C{i}",
2020,
"created",
(),
(),
f"[L{i}]",
"K",
i,
0,
"",
"fr",
True,
"",
)
for i in range(20) # 20 distinct codes, one priority row each
]
ev = LineageEvidence(
codes=tuple(f"C{i}" for i in range(20)),
families=(),
events=tuple(priorities),
element_diffs=(),
guidance=(),
max_prompt_rows=5,
)
rendered = ev.prompt_block()
kept = sum(rendered.count(f"[L{i}]") for i in range(20))
assert kept == 10 # 2 * max_prompt_rows(5)
class TestPayload:
def test_json_serializable_and_shaped(self):
@@ -744,9 +919,9 @@ class TestLineageEvidence:
(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 adopted.label == "CY2015 PFS final 79 FR 67716 ¶1251"
assert created.label == "CY2021 PFS final 85 FR 84547 ¶686"
assert replaces.label == "CY2021 PFS final 85 FR 84639 ¶1578"
assert replaces.from_codes == ("G2058",)
assert replaced_by.to_codes == ("99439",)
assert disappeared.label == "PFS CY2022 RVU file"
@@ -949,6 +1124,106 @@ class TestCachingAndScopedUrlResolution:
assert len(seen_keys) == len(set(seen_keys)) # never re-fetched
class TestConcurrentLineage:
"""C1/Ruling B14 — ``llm.lineage._store()`` must be thread-local:
bib.Store's sqlite connection is check_same_thread=True, and /chat
runs each turn in Starlette's threadpool. Mirrors
``TestConcurrentChats`` in test_evidence.py."""
@pytest.fixture
def bib_db(self, tmp_path):
from bib.item import Rule
from bib.store import Store
db = tmp_path / "bib.sqlite"
seed = Store(str(db), storage_dir=str(tmp_path / "storage"))
key = seed.create(
Rule(
title="Medicare Program; CY 2021 Payment Policies Under the PFS",
date_published="2020-12-28",
)
)
con = seed._con()
con.execute(
"INSERT INTO fr_anchor_docs "
"(item_key, document_number, html_url, start_page, end_page, fr_volume) "
"VALUES (?,?,?,?,?,?)",
(key, "2020-2021doc", "https://fr.test/2021-doc", 84000, 85000, 85),
)
con.execute(
"INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) "
"VALUES (?,?,?,?,?)",
(
key,
686,
84547,
1,
"We are finalizing HCPCS code G2058 as new CPT code 99439.",
),
)
con.commit()
seed.close()
return str(db), key
def test_no_cross_thread_sqlite_errors_one_store_per_thread(
self, bib_db, monkeypatch
):
from bib.store import Store
db_path, item_key = bib_db
opens: list[int] = []
real_init = Store.__init__
def counting_init(self, *a, **kw):
opens.append(threading.get_ident())
real_init(self, db_path)
monkeypatch.setattr(Store, "__init__", counting_init)
monkeypatch.setattr("conf.connect.bib", lambda: Store())
monkeypatch.setattr(
"bib.frlink.resolve",
lambda ref, *, store, item_key: SimpleNamespace(
url=f"https://fr.test/{item_key}#{ref}"
),
)
# I2 (later in this fix wave) adds "{vol} FR {page}" to this label
# once volume/page are resolvable — DE2VH9PD/YBM4IZUS-style fixture
# data does resolve them, so this constant moves with that change.
expected_label = "CY2021 PFS final 85 FR 84547 ¶686"
def event() -> LineageEvent:
store = lineage._store()
return lineage._to_lineage_event(
store, _ev("99439", 2021, "created", item_key=item_key, p_id=686), 1
)
errors: list[BaseException] = []
n_calls = 40
def worker(_i: int) -> None:
try:
for _ in range(n_calls):
store = lineage._store()
label = lineage._source_label(store, item_key, 686, 1)
url = lineage._fr_url(store, item_key, 686, 1)
assert label == expected_label, label
assert url == f"https://fr.test/{item_key}#p-686", url
ev = _lineage_evidence([event()])
(src,) = lineage.lineage_sources(store, ev, max_items=6)
assert src["label"] == expected_label
except BaseException as e: # noqa: BLE001 — collected, not raised in-thread
errors.append(e)
with ThreadPoolExecutor(max_workers=8) as pool:
list(pool.map(worker, range(8)))
assert errors == [], errors
# one Store construction per thread (threading.local caching),
# not one per call (320) and not one shared across every thread.
assert len(opens) == 8
assert len(set(opens)) == 8
def _lineage_evidence(
events, max_prompt_rows=25, explicit=(), element_diffs=()
) -> LineageEvidence:
@@ -1042,6 +1317,25 @@ class TestLineageSources:
out = lineage_sources(store, ev, max_items=2)
assert len(out) == 2
def test_event_loop_enforces_the_hard_cap_even_with_a_raised_max_items(self, store):
"""Ruling B11: _LINEAGE_SOURCES_HARD_CAP (10) is an absolute
ceiling the event loop must enforce itself — a config-raised
max_items must not let it through."""
events = []
for i in range(15):
key = f"MANY{i}"
_insert_extra_anchor(store, key, 2000 + i, 100 + i)
events.append(
lineage._to_lineage_event(
store,
_ev(f"M{i}", 2000 + i, "revalued", item_key=key, p_id=100 + i),
mtime=0,
)
)
ev = _lineage_evidence(events, max_prompt_rows=25)
out = lineage_sources(store, ev, max_items=20)
assert len(out) == 10
def test_label_equals_the_event_label(self, store):
e = lineage._to_lineage_event(
store,
@@ -1050,7 +1344,7 @@ class TestLineageSources:
)
ev = _lineage_evidence([e])
(out,) = lineage_sources(store, ev, max_items=6)
assert out["label"] == e.label == "CY2015 PFS final ¶1251"
assert out["label"] == e.label == "CY2015 PFS final 79 FR 67716 ¶1251"
assert out["id"] == e.label
assert out["kind"] == "rule"
assert out["item_key"] == "DE2VH9PD"
@@ -1227,10 +1521,12 @@ class TestLineageSourcesRulingB10:
assert [s["label"] for s in out] == [diffs[0].label, diffs[1].label]
def test_element_diffs_still_added_past_max_items_up_to_the_hard_cap(self, store):
# 8 distinct event paragraphs exactly fill max_items=8; two more
# distinct element-diff paragraphs are still appended (the "add
# up to 2 even when max_items has no room" rule), reaching the
# absolute hard cap of 10, not 11.
# 8 distinct event paragraphs (8 distinct codes, so Ruling B11's
# 4-per-code priority cap doesn't itself trim the selection)
# exactly fill max_items=8; two more distinct element-diff
# paragraphs are still appended (the "add up to 2 even when
# max_items has no room" rule), reaching the absolute hard cap
# of 10, not 11.
events = []
for i in range(8):
key = f"EV{i}"
@@ -1238,7 +1534,7 @@ class TestLineageSourcesRulingB10:
events.append(
lineage._to_lineage_event(
store,
_ev("Z", 2000 + i, "created", item_key=key, p_id=100 + i),
_ev(f"Z{i}", 2000 + i, "created", item_key=key, p_id=100 + i),
mtime=0,
)
)
@@ -1250,7 +1546,7 @@ class TestLineageSourcesRulingB10:
_diff("DX2", 501, "d2"),
_diff("DX3", 502, "d3"), # third diff — dropped by the element cap of 2
]
ev = _lineage_evidence(events, explicit=("Z",), element_diffs=diffs)
ev = _lineage_evidence(events, explicit=("Z0",), element_diffs=diffs)
out = lineage_sources(store, ev, max_items=8)
assert len(out) == 10
assert [s["label"] for s in out[8:]] == ["d1", "d2"]
@@ -1354,7 +1650,12 @@ class TestReconcileLabels:
class TestLineageEvidenceLive:
@pytest.mark.live
def test_median_under_300ms_on_live_replica(self):
def test_median_under_300ms_on_live_replica(self, restore_families):
# I5: lineage_evidence -> evidence._connect -> refresh_from(con)
# merges the live replica's derived pfs.code_family rows into the
# module-global FAMILIES registry — restore_families snapshots and
# restores it so this test doesn't leak thousands of derived
# families into every test that runs after it in the session.
from conf import ROOT
replica = ROOT / "data" / "replica" / "aco.ro.duckdb"

181
tests/llm/test_rag.py
View File

@@ -1,6 +1,7 @@
"""llm.rag — multi-collection retrieval + grounded streaming answer."""
import json
from dataclasses import replace
from datetime import date
from unittest.mock import MagicMock, patch
@@ -436,6 +437,21 @@ class TestEraBalance:
assert len(out) == 1
assert out[0].distance == 0.05
def test_coverage_when_eras_outnumber_top_n(self):
"""I3: 14 distinct eras (2013-2026), top_n=8 — without evenly
spacing the era selection, the round-robin's first pass alone
(14 eras, newest-first) fills the output and 2013 never
survives the final top_n slice. Both ends must survive."""
hits = [
self._rule_hit(f"R{year}", year, 0.1 + (2026 - year) * 0.01)
for year in range(2013, 2027)
]
out = era_balance(hits, per_era=2, top_n=8)
eras = {era_of(h) for h in out}
assert 2013 in eras
assert 2026 in eras
assert len(out) == 8
class TestBuildMessages:
def test_includes_labels_kinds_dates_and_rules(self):
@@ -574,6 +590,139 @@ class TestBuildMessages:
assert build_messages("q", []) == build_messages("q", [], lineage=None)
def _bmsrc(label, snippet="s", kind="corpus", date="2020-01-01"):
return {"label": label, "kind": kind, "date": date, "snippet": snippet}
def _val_row(code, year, label):
from pfs.valuation import ValuationRow
return ValuationRow(
code=code,
description="d",
vintage=f"CY{year} final",
year=year,
proposed=False,
status="A",
work=0.25,
pe_nf=0.22,
pe_f=0.06,
mp=0.02,
total_nf=0.49,
total_f=0.33,
cf=33.4009,
pay_nf=16.37,
pay_f=11.02,
label=label,
citation="90 FR 49266",
url="u",
)
class TestBuildMessagesBudget:
"""Ruling B11 — ``build_messages(..., budget_chars=...)`` drop order:
lineage-source excerpts (>4) -> retrieved (>6) -> cited (>8) ->
manual (>1) -> valuation rows (>12) -> lineage rows
(> lineage.max_prompt_rows). Each test proves the order with a
synthetic oversize input, matching against content built by
trimming the SAME way by hand (not a hand-counted char budget) so
the assertion stays correct if the rendered format ever changes."""
def test_budget_none_never_trims(self):
retrieved = [_bmsrc(f"R{i}") for i in range(50)]
msgs = build_messages("q", retrieved, n_retrieved=50, budget_chars=None)
assert msgs[1]["content"].count("[R") == 50
def test_under_budget_is_a_noop(self):
retrieved = [_bmsrc(f"R{i}") for i in range(3)]
msgs = build_messages("q", retrieved, n_retrieved=3, budget_chars=10_000_000)
assert msgs[1]["content"].count("[R") == 3
def test_lineage_source_excerpts_drop_first(self):
retrieved = [_bmsrc(f"R{i}") for i in range(10)]
ls = [_bmsrc(f"LS{i}", snippet="L" * 100) for i in range(8)]
sources = retrieved + ls
target = build_messages(
"q", retrieved + ls[:4], n_retrieved=10, n_lineage_sources=4
)[1]["content"]
msgs = build_messages(
"q",
sources,
n_retrieved=10,
n_lineage_sources=8,
budget_chars=len(target),
)
assert msgs[1]["content"] == target
assert msgs[1]["content"].count("[R") == 10 # retrieved untouched
def test_retrieved_drops_only_after_lineage_source_is_at_its_floor(self):
retrieved = [_bmsrc(f"R{i}", snippet="R" * 50) for i in range(10)]
ls = [_bmsrc(f"LS{i}", snippet="L" * 50) for i in range(8)]
cited = [_bmsrc(f"C{i}") for i in range(8)]
sources = retrieved + cited + ls
target = build_messages(
"q",
retrieved[:6] + cited + ls[:4],
n_retrieved=6,
n_cited=8,
n_lineage_sources=4,
)[1]["content"]
msgs = build_messages(
"q",
sources,
n_retrieved=10,
n_cited=8,
n_lineage_sources=8,
budget_chars=len(target),
)
assert msgs[1]["content"] == target
assert msgs[1]["content"].count("[C") == 8 # cited untouched — not its turn
def test_valuation_rows_capped_as_a_last_resort(self):
from llm.evidence import ValuationEvidence
rows = tuple(_val_row("G0556", 2010 + i, f"[Y{i}]") for i in range(15))
ev = ValuationEvidence(("G0556",), (), rows, (), max_prompt_rows=24)
big = build_messages("q", [], evidence=ev)[1]["content"]
msgs = build_messages("q", [], evidence=ev, budget_chars=len(big) - 1)
content = msgs[1]["content"]
assert sum(content.count(f"[Y{i}]") for i in range(15)) == 12
def test_lineage_rows_capped_as_the_final_resort(self):
events = tuple(
LineageEvent(
f"C{i}",
2020,
"revalued",
(),
(),
f"[L{i}]",
"K",
i,
0,
"",
"fr",
True,
"",
)
for i in range(20)
)
le = LineageEvidence(
codes=tuple(f"C{i}" for i in range(20)),
families=(),
events=events,
element_diffs=(),
guidance=(),
max_prompt_rows=20,
)
big = build_messages("q", [], lineage=le)[1]["content"]
msgs = build_messages("q", [], lineage=le, budget_chars=len(big) - 1)
content = msgs[1]["content"]
kept = sum(content.count(f"[L{i}]") for i in range(20))
assert kept <= 20 # budget-capped rendering never exceeds lineage_max_rows
assert kept < 20 # and it actually dropped something
class TestStreamAnswer:
def _pool(self, vram=24.0, serves_big=True):
pool = MagicMock()
@@ -961,7 +1110,6 @@ class TestStreamAnswer:
assert events[-2]["sources"] == [src, cited, manual]
mock_manual.assert_called_once_with(("APCM",), CFG)
@patch("llm.rag.detect_codes")
@patch("llm.rag._manual_sources", return_value=[])
@patch("llm.rag._engine")
@patch("llm.rag.code_cited_sources")
@@ -978,17 +1126,14 @@ class TestStreamAnswer:
mock_cited,
mock_engine,
mock_manual,
mock_detect,
):
from llm.evidence import ValuationEvidence
from pfs.families import Detection
mock_retrieve.return_value = []
mock_cited.return_value = []
mock_detect.return_value = Detection(
codes=(), families=("BIGFAM", "APCM"), explicit=(), wide=("BIGFAM",)
mock_ev.return_value = ValuationEvidence(
(), ("BIGFAM", "APCM"), (), (), wide=("BIGFAM",)
)
mock_ev.return_value = ValuationEvidence((), ("BIGFAM", "APCM"), (), ())
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}'])
@@ -1027,7 +1172,15 @@ class TestStreamAnswer:
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``)."""
``None``). A nonexistent replica path (fix-wave: this test used
to run against whatever real ``data/replica/aco.ro.duckdb`` the
checkout happens to have — ``warm()`` opens it unconditionally
before the code check, silently merging real derived families
into the process-global registry and leaking them into every
test that runs afterward in the same session) — the byte-
identical assertion below is exactly as meaningful either way,
since a codeless question never reaches a query regardless."""
cfg = replace(CFG, duckdb_replica="/nonexistent/aco.ro.duckdb")
mock_retrieve.return_value = []
client = MockClient.return_value.__enter__.return_value
resp = client.stream.return_value.__enter__.return_value
@@ -1038,7 +1191,7 @@ class TestStreamAnswer:
)
return list(
stream_answer(
"why did CMS finalize this policy?", cfg=CFG, pool=self._pool()
"why did CMS finalize this policy?", cfg=cfg, pool=self._pool()
)
)
@@ -1079,3 +1232,15 @@ class TestManualSourcesWrapper:
def test_unopenable_replica_yields_empty(self, _connect, caplog):
assert rag._manual_sources(("CCM",), CFG) == []
assert "manual sources skipped" in caplog.text
@patch("llm.rag.manual_sources", side_effect=RuntimeError("boom"))
@patch("llm.rag._connect")
@patch("llm.rag._replica_path", return_value="/x/aco.ro.duckdb")
def test_query_failure_yields_empty_and_still_closes_the_cursor(
self, _path, mock_connect, _manual, caplog
):
cur = MagicMock()
mock_connect.return_value.cursor.return_value = cur
assert rag._manual_sources(("CCM",), CFG) == []
assert "manual sources skipped" in caplog.text
cur.close.assert_called_once()

18
tests/pfs/test_families.py
View File

@@ -38,23 +38,7 @@ from pfs.families import (
stem_tokens,
)
@pytest.fixture
def restore_families():
"""Snapshot ``pfs.families.FAMILIES`` and restore it (and the
compiled phrase/code index built off it) in teardown, even if the
test body raises — a test that calls ``refresh_from``/mutates
``FAMILIES`` directly must not leave the live registry or its index
clobbered for the rest of the session."""
from pfs import families as mod
before = dict(mod.FAMILIES)
try:
yield mod
finally:
mod.FAMILIES.clear()
mod.FAMILIES.update(before)
mod.rebuild_index()
# restore_families fixture: tests/conftest.py (shared with tests/llm — I5).
class TestFindCodes: