feat(llm): era-balanced retrieval for history questions — retrieve(mode="timeline"), auto trigger, ChatRequest.mode (refs #691)

This commit is contained in:
kert
2026-09-10 10:54:55 -04:00
parent c0ba1a9b69
commit 12129d8a6e
6 changed files with 530 additions and 22 deletions

View File

@@ -44,6 +44,7 @@ app = FastAPI(
) )
_ISO_DATE = re.compile(r"^\d{4}-\d{2}-\d{2}$") _ISO_DATE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
_MODES = {"auto", "timeline", "recent"}
# OTel instrumentation — no-op if perf not installed or telemetry disabled. # OTel instrumentation — no-op if perf not installed or telemetry disabled.
try: try:
@@ -59,6 +60,13 @@ except ImportError: # pragma: no cover — perf always installed in the image
class ChatRequest(BaseModel): class ChatRequest(BaseModel):
question: str question: str
since: str | None = None since: str | None = None
#: "auto" (default — era-balanced retrieval for history-shaped
#: questions, otherwise recency-blended), "timeline" or "recent" —
#: forced from the UI (Task 6) and the golden eval (Task 7).
#: Validated by ``chat()`` (400, not pydantic's 422) against
#: ``_MODES``; the type stays ``str`` so an invalid value reaches
#: that check instead of failing model validation first.
mode: str = "auto"
def _page() -> str: def _page() -> str:
@@ -82,7 +90,7 @@ def whoami(x_auth_request_user: str = Header(default="")) -> dict:
return {"user": x_auth_request_user} return {"user": x_auth_request_user}
def _sse(question: str, since: str = "") -> Iterator[str]: def _sse(question: str, since: str = "", mode: str = "auto") -> Iterator[str]:
from llm import config as llm_config from llm import config as llm_config
from llm.pool import HostPool from llm.pool import HostPool
from llm.rag import stream_answer from llm.rag import stream_answer
@@ -90,7 +98,9 @@ def _sse(question: str, since: str = "") -> Iterator[str]:
cfg = llm_config.load() cfg = llm_config.load()
pool = HostPool.from_config(cfg) pool = HostPool.from_config(cfg)
try: try:
for event in stream_answer(question, cfg=cfg, pool=pool, since=since): for event in stream_answer(
question, cfg=cfg, pool=pool, since=since, mode=mode
):
yield f"data: {json.dumps(event)}\n\n" yield f"data: {json.dumps(event)}\n\n"
except Exception as exc: # surface to the transcript, don't 500 mid-stream except Exception as exc: # surface to the transcript, don't 500 mid-stream
yield f"data: {json.dumps({'type': 'error', 'message': str(exc)})}\n\n" yield f"data: {json.dumps({'type': 'error', 'message': str(exc)})}\n\n"
@@ -104,7 +114,14 @@ def chat(req: ChatRequest) -> StreamingResponse:
since = (req.since or "").strip() since = (req.since or "").strip()
if since and not _ISO_DATE.match(since): if since and not _ISO_DATE.match(since):
raise HTTPException(status_code=400, detail="since must be YYYY-MM-DD") raise HTTPException(status_code=400, detail="since must be YYYY-MM-DD")
return StreamingResponse(_sse(question, since), media_type="text/event-stream") mode = (req.mode or "auto").strip()
if mode not in _MODES:
raise HTTPException(
status_code=400, detail="mode must be one of: auto, timeline, recent"
)
return StreamingResponse(
_sse(question, since, mode), media_type="text/event-stream"
)
@app.get("/hosts") @app.get("/hosts")

View File

@@ -44,6 +44,8 @@ class LlmConfig:
default_factory=lambda: dict(_DEFAULT_K_PER_KIND) default_factory=lambda: dict(_DEFAULT_K_PER_KIND)
) )
top_n: int = 8 top_n: int = 8
timeline_per_era: int = 2 # era_balance: hits kept per rule year on the first pass
timeline_overfetch: int = 6 # era_balance: k multiplier per kind (vs _OVERFETCH=3) so every era has candidates
duckdb_replica: str = "data/replica/aco.ro.duckdb" duckdb_replica: str = "data/replica/aco.ro.duckdb"
valuation_years: int = 4 valuation_years: int = 4
code_cited_per_code: int = 3 code_cited_per_code: int = 3
@@ -114,6 +116,8 @@ def load() -> LlmConfig:
recency_weight=float(_opt(section, "recency_weight", 0.3)), recency_weight=float(_opt(section, "recency_weight", 0.3)),
k_per_kind=k_per_kind, k_per_kind=k_per_kind,
top_n=int(_opt(section, "top_n", 8)), top_n=int(_opt(section, "top_n", 8)),
timeline_per_era=int(_opt(section, "timeline_per_era", 2)),
timeline_overfetch=int(_opt(section, "timeline_overfetch", 6)),
duckdb_replica=os.environ.get("LLM_DUCKDB_REPLICA") duckdb_replica=os.environ.get("LLM_DUCKDB_REPLICA")
or str(_opt(section, "duckdb_replica", "data/replica/aco.ro.duckdb")), or str(_opt(section, "duckdb_replica", "data/replica/aco.ro.duckdb")),
valuation_years=int(_opt(section, "valuation_years", 4)), valuation_years=int(_opt(section, "valuation_years", 4)),

View File

@@ -5,7 +5,12 @@ Single-shot and stateless — each question is retrieved and answered on
its own. Retrieval embeds the question once and searches the its own. Retrieval embeds the question once and searches the
``comments``, ``rules`` and ``corpus`` pgvector collections; hits are ``comments``, ``rules`` and ``corpus`` pgvector collections; hits are
merged and re-ranked by similarity × recency (``llm.rerank``), then each merged and re-ranked by similarity × recency (``llm.rerank``), then each
source gets a deep link (``llm.links``). Generation streams from the source gets a deep link (``llm.links``). A history-shaped question
("history of CCM", two-or-more-years, "when did…") instead gets
era-balanced retrieval (``retrieve(mode="timeline")`` /
``era_balance``): similarity-only scoring, spread across rule years
rather than weighted toward the newest — see ``is_history_question``.
Generation streams from the
largest live Ollama host (``HostPool.acquire_generation``) with the largest live Ollama host (``HostPool.acquire_generation``) with the
model tier that host can hold (``pick_model``). When the question model tier that host can hold (``pick_model``). When the question
cites HCPCS/CPT codes, ``llm.evidence``/``llm.lineage`` supply cites HCPCS/CPT codes, ``llm.evidence``/``llm.lineage`` supply
@@ -18,6 +23,10 @@ SSE event order is ``lineage`` (only when detected) → ``valuation``
from __future__ import annotations from __future__ import annotations
import json import json
import os
import re
import threading
from dataclasses import replace
from datetime import date from datetime import date
from typing import Iterator from typing import Iterator
@@ -33,14 +42,166 @@ from llm.evidence import (
) )
from llm.index import _engine from llm.index import _engine
from llm.lineage import LineageEvidence from llm.lineage import LineageEvidence
from llm.lineage import _store as _bib_store
from llm.links import as_source from llm.links import as_source
from llm.pool import HostPool, PoolEmbeddings, pick_model from llm.pool import HostPool, PoolEmbeddings, pick_model
from llm.rerank import Hit, blend, filter_since from llm.rerank import Hit, _similarity, blend, filter_since
from pfs.descriptors import rule_year_of
_TIMEOUT = httpx.Timeout(300.0, connect=5.0) _TIMEOUT = httpx.Timeout(300.0, connect=5.0)
_COLLECTIONS = {"comment": "comments", "rule": "rules", "corpus": "corpus"} _COLLECTIONS = {"comment": "comments", "rule": "rules", "corpus": "corpus"}
_OVERFETCH = 3 _OVERFETCH = 3
#: History/timeline-shaped questions — "history of CCM", "what replaced
#: G2058", "when was 99490 created" — trigger era-balanced retrieval
#: instead of the recency blend (``retrieve(mode="auto")``).
HISTORY_RE = re.compile(
r"\b("
r"history|histories|historical|evolution|evolved|"
r"over the years|over time|timeline|"
r"when (?:did|was|were)|"
r"replaced|replacement|predecessor|successor|"
r"since (?:19|20)\d\d|"
r"(?:from|between) (?:19|20)\d\d"
r")\b",
re.IGNORECASE,
)
_YEAR_RE = re.compile(r"\b(?:19|20)\d\d\b")
def is_history_question(question: str) -> bool:
"""True for a question shaped like "history of CCM" or one naming
two or more distinct years — either is a cue that the answer spans
multiple rule eras, so retrieval should balance across them rather
than lean on the most recent material."""
q = question or ""
if HISTORY_RE.search(q):
return True
return len(set(_YEAR_RE.findall(q))) >= 2
# ── era_of: the rule year a hit belongs to, for era-balanced retrieval ──
#: docket id -> PFS rule year, cached in-process and rebuilt only when
#: the bib sqlite file's mtime changes (or on first use) — ``era_of``
#: is only ever asked about a comment hit's docket, so a question with
#: no comment hits never triggers ``Store.dockets()`` at all.
_UNSET = object()
_docket_era_mtime: object = _UNSET
_docket_era_map: dict[str, int] = {}
_docket_era_lock = threading.Lock()
def _docket_rule_year(comment_end_date: str) -> int:
"""A PFS docket's comment period closes in the second half of a
calendar year (Sept, typically) for a rule governing the NEXT
calendar year; one closing in the first half belongs to that same
year's rule. 0 when unparsable."""
try:
year = int(comment_end_date[:4])
month = int(comment_end_date[5:7])
except (ValueError, TypeError, IndexError):
return 0
return year + 1 if month >= 7 else year
def _docket_year(docket_id: str) -> int:
"""The PFS rule year *docket_id* belongs to, 0 when unknown."""
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)
def _year_of(iso_date: str) -> int:
try:
return int((iso_date or "")[:4])
except ValueError:
return 0
def era_of(hit: Hit) -> int:
"""The rule year *hit* belongs to: ``rule_year_of`` for rules, the
comment's docket's rule year (falling back to the comment's own
date) for comments, the hit's date year for corpus material; 0 —
its own era, sorted last — when unknown."""
kind = hit.metadata.get("kind", "")
date_str = hit.metadata.get("date", "")
if kind == "rule":
return rule_year_of(hit.metadata.get("title", ""), date_str)
if kind == "comment":
year = _docket_year(hit.metadata.get("docket", ""))
return year if year else _year_of(date_str)
return _year_of(date_str)
def era_balance(hits: list[Hit], *, per_era: int, top_n: int) -> list[Hit]:
"""Era-balanced re-ranking for history questions: score every hit
by similarity alone (no recency term), dedupe per ``item_key``
(best chunk wins, exactly like ``blend``), then take the best
``per_era`` hits per era and interleave round-robin across eras —
newest first, one hit per era before a second — filling any
remaining slots to ``top_n`` from the rest by score. Deterministic:
ties broken by ``item_key``."""
best: dict[str, Hit] = {}
for h in hits:
scored = replace(h, score=_similarity(h.distance))
key = h.metadata.get("item_key", "") or str(id(h))
if key not in best or scored.score > best[key].score:
best[key] = scored
def _tiebreak(h: Hit) -> str:
md = h.metadata
return md.get("item_key", "") or md.get("p_id", "") or md.get("seq", "")
ranked = sorted(best.values(), key=lambda h: (-h.score, _tiebreak(h)))
by_era: dict[int, list[Hit]] = {}
for h in ranked:
by_era.setdefault(era_of(h), []).append(h)
eras_desc = sorted(by_era, reverse=True)
capped = {era: by_era[era][:per_era] for era in eras_desc}
selected: list[Hit] = []
for round_idx in range(per_era):
for era in eras_desc:
bucket = capped[era]
if round_idx < len(bucket):
selected.append(bucket[round_idx])
chosen = {id(h) for h in selected}
remainder = [h for h in ranked if id(h) not in chosen]
if len(selected) < top_n:
selected.extend(remainder[: top_n - len(selected)])
return selected[:top_n]
def _resolve_mode(question: str, mode: str) -> str:
"""``mode`` as-is for "timeline"/"recent"; "auto" resolves to
"timeline" for a history-shaped question, else "recent"."""
if mode == "timeline":
return "timeline"
if mode == "auto":
return "timeline" if is_history_question(question) else "recent"
return "recent"
_SYSTEM = ( _SYSTEM = (
"You answer questions about CMS rulemaking using ONLY the excerpts " "You answer questions about CMS rulemaking using ONLY the excerpts "
"provided below. Excerpts come from three kinds of sources: public " "provided below. Excerpts come from three kinds of sources: public "
@@ -68,7 +229,13 @@ _SYSTEM = (
) )
def _hits(question_vec: list[float], *, cfg: LlmConfig, pool: HostPool) -> list[Hit]: def _hits(
question_vec: list[float],
*,
cfg: LlmConfig,
pool: HostPool,
overfetch: int = _OVERFETCH,
) -> list[Hit]:
from llm.index import vectorstore from llm.index import vectorstore
hits: list[Hit] = [] hits: list[Hit] = []
@@ -78,7 +245,7 @@ def _hits(question_vec: list[float], *, cfg: LlmConfig, pool: HostPool) -> list[
continue continue
store = vectorstore(collection, cfg, pool) store = vectorstore(collection, cfg, pool)
for doc, distance in store.similarity_search_with_score_by_vector( for doc, distance in store.similarity_search_with_score_by_vector(
question_vec, k=k * _OVERFETCH question_vec, k=k * overfetch
): ):
md = {key: str(v) for key, v in (doc.metadata or {}).items()} md = {key: str(v) for key, v in (doc.metadata or {}).items()}
md.setdefault("kind", kind) md.setdefault("kind", kind)
@@ -99,12 +266,27 @@ def retrieve(
pool: HostPool, pool: HostPool,
since: str = "", since: str = "",
now: date | None = None, now: date | None = None,
mode: str = "auto",
) -> list[dict]: ) -> list[dict]:
"""Top sources for ``question`` across all collections, recency-blended. """Top sources for ``question`` across all collections.
``since`` (ISO date) hard-filters to material dated on/after it. ``since`` (ISO date) hard-filters to material dated on/after it.
``mode``: "recent" (default behaviour — similarity × recency blend,
byte-identical to before Task 4) or "timeline" (era-balanced —
overfetches ``cfg.timeline_overfetch`` per kind and skips the
recency blend in favour of ``era_balance``, so a history question
sees material spread across rule years instead of mostly the
newest). "auto" picks "timeline" iff ``is_history_question(question)``.
""" """
resolved = _resolve_mode(question, mode)
vec = PoolEmbeddings(pool, cfg.embed_model).embed_query(question) vec = PoolEmbeddings(pool, cfg.embed_model).embed_query(question)
if resolved == "timeline":
hits = filter_since(
_hits(vec, cfg=cfg, pool=pool, overfetch=cfg.timeline_overfetch), since
)
ranked = era_balance(hits, per_era=cfg.timeline_per_era, top_n=cfg.top_n)
else:
hits = filter_since(_hits(vec, cfg=cfg, pool=pool), since) hits = filter_since(_hits(vec, cfg=cfg, pool=pool), since)
ranked = blend( ranked = blend(
hits, hits,
@@ -147,10 +329,19 @@ def build_messages(
def stream_answer( def stream_answer(
question: str, *, cfg: LlmConfig, pool: HostPool, since: str = "" question: str,
*,
cfg: LlmConfig,
pool: HostPool,
since: str = "",
mode: str = "auto",
) -> Iterator[dict]: ) -> Iterator[dict]:
"""Retrieve, then stream a grounded answer from the largest live host. """Retrieve, then stream a grounded answer from the largest live host.
``mode`` ("auto"/"timeline"/"recent") is forwarded to ``retrieve``;
"auto" (the default) switches to era-balanced retrieval on its own
for a history-shaped question (``is_history_question``).
When the question cites HCPCS/CPT codes, yields a ``lineage`` event When the question cites HCPCS/CPT codes, yields a ``lineage`` event
(evidence's ``payload()``, the dated timeline) then a ``valuation`` (evidence's ``payload()``, the dated timeline) then a ``valuation``
event (RVUs and payment) — both before any tokens — and merges event (RVUs and payment) — both before any tokens — and merges
@@ -158,9 +349,11 @@ def stream_answer(
or, for a detected family, its family key — into ``sources``. Then or, for a detected family, its family key — into ``sources``. Then
yields ``{"type":"token","text":…}`` events as the model generates, yields ``{"type":"token","text":…}`` events as the model generates,
then one ``{"type":"sources", "sources": […], "model": …, "host": then one ``{"type":"sources", "sources": […], "model": …, "host":
…}`` and a final ``{"type":"done"}``. …, "mode": "timeline"|"recent"}`` (the resolved mode) and a final
``{"type":"done"}``.
""" """
sources = retrieve(question, cfg=cfg, pool=pool, since=since) resolved_mode = _resolve_mode(question, mode)
sources = retrieve(question, cfg=cfg, pool=pool, since=since, mode=mode)
lineage = lineage_evidence(question, cfg) lineage = lineage_evidence(question, cfg)
evidence = valuation_evidence(question, cfg) evidence = valuation_evidence(question, cfg)
codes: tuple[str, ...] = evidence.codes if evidence is not None else () codes: tuple[str, ...] = evidence.codes if evidence is not None else ()
@@ -206,5 +399,11 @@ def stream_answer(
yield {"type": "token", "text": chunk} yield {"type": "token", "text": chunk}
if data.get("done"): if data.get("done"):
break break
yield {"type": "sources", "sources": sources, "model": model, "host": host} yield {
"type": "sources",
"sources": sources,
"model": model,
"host": host,
"mode": resolved_mode,
}
yield {"type": "done"} yield {"type": "done"}

View File

@@ -124,6 +124,8 @@ pg_user = "llm"
recency_half_life_days = 365 recency_half_life_days = 365
recency_weight = 0.3 recency_weight = 0.3
top_n = 8 top_n = 8
timeline_per_era = 2 # era_balance: hits kept per rule year on the first pass (mode="timeline")
timeline_overfetch = 6 # era_balance: k multiplier per kind for timeline mode (vs the recent-mode ×3)
duckdb_replica = "data/replica/aco.ro.duckdb" # read-only DuckDB replica the chat reads valuations from duckdb_replica = "data/replica/aco.ro.duckdb" # read-only DuckDB replica the chat reads valuations from
valuation_years = 4 # final-rule vintages shown per code (plus the newest NPRM) valuation_years = 4 # final-rule vintages shown per code (plus the newest NPRM)
code_cited_per_code = 3 # excerpts literally citing each detected code (or family) code_cited_per_code = 3 # excerpts literally citing each detected code (or family)

View File

@@ -94,6 +94,26 @@ class TestChatSince:
assert r.status_code == 400 assert r.status_code == 400
class TestChatMode:
@patch("llm.rag.stream_answer")
def test_mode_forwarded(self, mock_stream):
mock_stream.return_value = iter([{"type": "done"}])
r = client.post("/chat", json={"question": "q", "mode": "timeline"})
assert r.status_code == 200
assert mock_stream.call_args.kwargs["mode"] == "timeline"
@patch("llm.rag.stream_answer")
def test_mode_defaults_to_auto(self, mock_stream):
mock_stream.return_value = iter([{"type": "done"}])
r = client.post("/chat", json={"question": "q"})
assert r.status_code == 200
assert mock_stream.call_args.kwargs["mode"] == "auto"
def test_bad_mode_400(self):
r = client.post("/chat", json={"question": "q", "mode": "x"})
assert r.status_code == 400
def _cfg(**kw): def _cfg(**kw):
from llm.config import LlmConfig from llm.config import LlmConfig

View File

@@ -7,9 +7,18 @@ from unittest.mock import MagicMock, patch
import pytest import pytest
from langchain_core.documents import Document from langchain_core.documents import Document
import llm.rag as rag
from llm.config import LlmConfig from llm.config import LlmConfig
from llm.lineage import LineageEvent, LineageEvidence from llm.lineage import LineageEvent, LineageEvidence
from llm.rag import build_messages, retrieve, stream_answer from llm.rag import (
build_messages,
era_balance,
era_of,
is_history_question,
retrieve,
stream_answer,
)
from llm.rerank import Hit
CFG = LlmConfig( CFG = LlmConfig(
ollama_hosts=("http://h1:11434",), ollama_hosts=("http://h1:11434",),
@@ -29,6 +38,21 @@ CFG = LlmConfig(
NOW = date(2026, 9, 3) NOW = date(2026, 9, 3)
@pytest.fixture(autouse=True)
def _reset_docket_era_cache():
"""``era_of``'s comment→docket-year map is a module-global cache
(llm.rag mirrors lineage.py's ``_STORE``/``_ReplicaCache`` pattern)
— no test may inherit another's mocked store or mapping."""
def _reset():
rag._docket_era_mtime = rag._UNSET
rag._docket_era_map.clear()
_reset()
yield
_reset()
def _doc(text, **md): def _doc(text, **md):
return Document(page_content=text, metadata=md) return Document(page_content=text, metadata=md)
@@ -203,6 +227,215 @@ class TestRetrieve:
(s,) = retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW) (s,) = retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
assert s["kind"] == "comment" and s["label"] == "C-9" assert s["kind"] == "comment" and s["label"] == "C-9"
@patch("llm.rag.era_balance")
@patch("llm.rag.blend")
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_timeline_mode_overfetches_and_skips_blend(
self, mock_vs, MockEmb, mock_blend, mock_balance
):
MockEmb.return_value.embed_query.return_value = [0.0]
stores = {}
def factory(collection, cfg, pool):
s = MagicMock()
s.similarity_search_with_score_by_vector.return_value = []
stores[collection] = s
return s
mock_vs.side_effect = factory
mock_balance.return_value = []
retrieve("history of CCM", cfg=CFG, pool=MagicMock(), mode="timeline")
comments = stores["comments"].similarity_search_with_score_by_vector
rules = stores["rules"].similarity_search_with_score_by_vector
# cfg.timeline_overfetch defaults to 6, vs _OVERFETCH=3 for "recent"
assert comments.call_args.kwargs["k"] == 2 * 6
assert rules.call_args.kwargs["k"] == 1 * 6
mock_blend.assert_not_called()
mock_balance.assert_called_once()
assert mock_balance.call_args.kwargs == {
"per_era": CFG.timeline_per_era,
"top_n": CFG.top_n,
}
@patch("llm.rag.era_balance")
@patch("llm.rag.blend")
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_mode_recent_forces_blend_even_for_a_history_question(
self, mock_vs, MockEmb, mock_blend, mock_balance
):
MockEmb.return_value.embed_query.return_value = [0.0]
mock_vs.side_effect = _stores({})
mock_blend.return_value = []
retrieve("history of CCM", cfg=CFG, pool=MagicMock(), mode="recent", now=NOW)
mock_blend.assert_called_once()
mock_balance.assert_not_called()
@patch("llm.rag.era_balance")
@patch("llm.rag.blend")
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_mode_timeline_forces_era_balance_for_a_plain_question(
self, mock_vs, MockEmb, mock_blend, mock_balance
):
MockEmb.return_value.embed_query.return_value = [0.0]
mock_vs.side_effect = _stores({})
mock_balance.return_value = []
retrieve("telehealth", cfg=CFG, pool=MagicMock(), mode="timeline")
mock_balance.assert_called_once()
mock_blend.assert_not_called()
@patch("llm.rag.era_balance")
@patch("llm.rag.blend")
@patch("llm.rag.PoolEmbeddings")
@patch("llm.index.vectorstore")
def test_mode_auto_routes_to_timeline_for_a_history_question(
self, mock_vs, MockEmb, mock_blend, mock_balance
):
MockEmb.return_value.embed_query.return_value = [0.0]
mock_vs.side_effect = _stores({})
mock_balance.return_value = []
retrieve("history of CCM", cfg=CFG, pool=MagicMock()) # mode="auto" default
mock_balance.assert_called_once()
mock_blend.assert_not_called()
class TestIsHistoryQuestion:
@pytest.mark.parametrize(
"question",
[
"history of CCM",
"what replaced G2058",
"how did payment change between 2015 and 2025",
"when was 99490 created",
],
)
def test_positives(self, question):
assert is_history_question(question) is True
@pytest.mark.parametrize(
"question",
[
"what is the 2026 payment for 99490",
"does 99490 pass telehealth step 3",
],
)
def test_negatives(self, question):
assert is_history_question(question) is False
def test_two_distinct_years_without_a_trigger_word_counts_as_history(self):
assert is_history_question("99490 2015 vs 99490 2027 payment") is True
def test_one_year_alone_is_not_enough(self):
assert is_history_question("the 2026 payment for 99490") is False
class TestEraOf:
def _hit(self, **md):
return Hit(text="t", metadata={k: str(v) for k, v in md.items()}, distance=0.1)
def test_rule_uses_rule_year_of(self):
h = self._hit(kind="rule", title="CY2021 PFS final", date="2020-11-15")
assert era_of(h) == 2021
def test_corpus_uses_date_year(self):
h = self._hit(kind="corpus", date="2019-03-01")
assert era_of(h) == 2019
def test_corpus_unknown_date_is_era_zero(self):
h = self._hit(kind="corpus", date="")
assert era_of(h) == 0
@patch("llm.rag._bib_store")
def test_comment_docket_closing_in_september_belongs_to_next_year(self, mock_store):
from bib.dockets import Docket
mock_store.return_value.dockets.return_value = [
Docket(id="CMS-2023-9999", comment_end_date="2023-09-15")
]
mock_store.return_value._db_path = "/nonexistent/bib.sqlite"
h = self._hit(
kind="comment", docket="CMS-2023-9999", date="2023-09-10", item_key="K"
)
assert era_of(h) == 2024
@patch("llm.rag._bib_store")
def test_comment_docket_closing_in_january_belongs_to_that_year(self, mock_store):
from bib.dockets import Docket
mock_store.return_value.dockets.return_value = [
Docket(id="CMS-2024-1", comment_end_date="2024-01-10")
]
mock_store.return_value._db_path = "/nonexistent/bib.sqlite"
h = self._hit(
kind="comment", docket="CMS-2024-1", date="2024-01-05", item_key="K"
)
assert era_of(h) == 2024
@patch("llm.rag._bib_store")
def test_comment_unknown_docket_falls_back_to_its_own_date(self, mock_store):
mock_store.return_value.dockets.return_value = []
mock_store.return_value._db_path = "/nonexistent/bib.sqlite"
h = self._hit(kind="comment", docket="CMS-9999-1", date="2019-05-01")
assert era_of(h) == 2019
@patch("llm.rag._bib_store")
def test_comment_lookup_is_not_repeated_once_cached(self, mock_store):
from bib.dockets import Docket
mock_store.return_value.dockets.return_value = [
Docket(id="CMS-2024-1", comment_end_date="2024-01-10")
]
mock_store.return_value._db_path = "/nonexistent/bib.sqlite"
h1 = self._hit(kind="comment", docket="CMS-2024-1", date="2024-01-05")
h2 = self._hit(kind="comment", docket="CMS-2024-1", date="2024-01-05")
assert era_of(h1) == era_of(h2) == 2024
assert mock_store.return_value.dockets.call_count == 1
class TestEraBalance:
def _rule_hit(self, key, year, distance):
return Hit(
text="t",
metadata={
"kind": "rule",
"title": f"CY{year} PFS final",
"date": f"{year - 1}-11-01",
"item_key": key,
},
distance=distance,
)
def test_picks_older_and_newer_eras_over_higher_scoring_middle_era(self):
hits = [
self._rule_hit("R1", 2026, 0.05), # score .95 — best 2026
self._rule_hit("R2", 2026, 0.10), # score .90
self._rule_hit("R3", 2026, 0.15), # score .85
self._rule_hit("R4", 2015, 0.40), # score .60 — only CY2015 hit
self._rule_hit("R5", 2027, 0.35), # score .65 — only CY2027 hit
]
out = era_balance(hits, per_era=1, top_n=3)
assert [h.metadata["item_key"] for h in out] == ["R5", "R1", "R4"]
def test_deterministic_tiebreak_by_item_key(self):
hits = [
self._rule_hit("Z", 2026, 0.10),
self._rule_hit("A", 2026, 0.10),
]
first = [h.metadata["item_key"] for h in era_balance(hits, per_era=2, top_n=2)]
again = [h.metadata["item_key"] for h in era_balance(hits, per_era=2, top_n=2)]
assert first == again == ["A", "Z"]
def test_dedupes_per_item_key_keeping_best_chunk(self):
hits = [
self._rule_hit("A", 2026, 0.30),
self._rule_hit("A", 2026, 0.05), # better chunk of the same item
]
out = era_balance(hits, per_era=2, top_n=2)
assert len(out) == 1
assert out[0].distance == 0.05
class TestBuildMessages: class TestBuildMessages:
def test_includes_labels_kinds_dates_and_rules(self): def test_includes_labels_kinds_dates_and_rules(self):
@@ -381,7 +614,9 @@ class TestStreamAnswer:
pool.check.assert_called_once_with("chat") pool.check.assert_called_once_with("chat")
# no codes in the question — pgvector is never asked for cited rules # no codes in the question — pgvector is never asked for cited rules
mock_engine.assert_not_called() mock_engine.assert_not_called()
mock_retrieve.assert_called_once_with("q", cfg=CFG, pool=pool, since="") mock_retrieve.assert_called_once_with(
"q", cfg=CFG, pool=pool, since="", mode="auto"
)
assert events[0] == {"type": "token", "text": "Doc"} assert events[0] == {"type": "token", "text": "Doc"}
assert events[1] == {"type": "token", "text": "tors"} assert events[1] == {"type": "token", "text": "tors"}
assert events[-2] == { assert events[-2] == {
@@ -389,6 +624,7 @@ class TestStreamAnswer:
"sources": [src], "sources": [src],
"model": "big", "model": "big",
"host": "http://h1:11434", "host": "http://h1:11434",
"mode": "recent",
} }
assert events[-1] == {"type": "done"} assert events[-1] == {"type": "done"}
body = client.stream.call_args.kwargs["json"] body = client.stream.call_args.kwargs["json"]
@@ -421,6 +657,36 @@ class TestStreamAnswer:
list(stream_answer("q", cfg=CFG, pool=self._pool(), since="2025-09-01")) list(stream_answer("q", cfg=CFG, pool=self._pool(), since="2025-09-01"))
assert mock_retrieve.call_args.kwargs["since"] == "2025-09-01" 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_mode_forwarded_to_retrieve_and_resolved_onto_sources_event(
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
resp.iter_lines.return_value = iter(['{"message":{"content":""},"done":true}'])
events = list(stream_answer("q", cfg=CFG, pool=self._pool(), mode="timeline"))
assert mock_retrieve.call_args.kwargs["mode"] == "timeline"
assert events[-2]["mode"] == "timeline"
@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_mode_auto_resolves_to_timeline_for_a_history_question(
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
resp.iter_lines.return_value = iter(['{"message":{"content":""},"done":true}'])
events = list(stream_answer("history of CCM", cfg=CFG, pool=self._pool()))
assert mock_retrieve.call_args.kwargs["mode"] == "auto"
assert events[-2]["mode"] == "timeline"
@patch("llm.rag.valuation_evidence", return_value=None) @patch("llm.rag.valuation_evidence", return_value=None)
@patch("llm.rag.lineage_evidence", return_value=None) @patch("llm.rag.lineage_evidence", return_value=None)
@patch("llm.rag.httpx.Client") @patch("llm.rag.httpx.Client")