feat(llm): era-balanced retrieval for history questions — retrieve(mode="timeline"), auto trigger, ChatRequest.mode (refs #691)
This commit is contained in:
@@ -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")
|
||||||
|
|||||||
@@ -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)),
|
||||||
|
|||||||
217
src/llm/rag.py
217
src/llm/rag.py
@@ -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"}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
Reference in New Issue
Block a user