0cf878d4ab
Some checks failed
CI / lint (push) Successful in 31s
CI / notebooks-smoke (push) Successful in 1m25s
Deploy / notebooks (push) Has been skipped
Deploy / zotero (push) Has been skipped
Deploy / docs (push) Has been skipped
Deploy / api (push) Has been skipped
Deploy / llm (push) Has been skipped
Deploy / mc (push) Has been skipped
Infra CI / notebooks (push) Successful in 1m5s
Infra CI / zotero (push) Successful in 13s
Infra CI / docs (push) Successful in 27s
Infra CI / api (push) Successful in 58s
Infra CI / llm (push) Successful in 40s
Infra CI / mc (push) Failing after 13s
Deploy / report (push) Successful in 18s
CI / test (push) Successful in 14m10s
1451 lines
53 KiB
Python
1451 lines
53 KiB
Python
"""llm.rag — multi-collection retrieval + grounded streaming answer."""
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import replace
|
|
from datetime import date
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from langchain_core.documents import Document
|
|
|
|
import llm.rag as rag
|
|
from llm.config import LlmConfig
|
|
from llm.lineage import LineageEvent, LineageEvidence
|
|
from llm.rag import (
|
|
build_messages,
|
|
era_balance,
|
|
era_of,
|
|
is_history_question,
|
|
retrieve,
|
|
stream_answer,
|
|
)
|
|
from llm.rerank import Hit
|
|
|
|
CFG = LlmConfig(
|
|
ollama_hosts=("http://h1:11434",),
|
|
host_vram={"http://h1:11434": 24},
|
|
embed_model="embed",
|
|
instruct_model="chat",
|
|
instruct_model_large="big",
|
|
embed_dim=768,
|
|
pg_host="x",
|
|
pg_port=5432,
|
|
pg_db="llm",
|
|
pg_user="llm",
|
|
build_ann_index=True,
|
|
k_per_kind={"comment": 2, "rule": 1, "corpus": 1},
|
|
top_n=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):
|
|
return Document(page_content=text, metadata=md)
|
|
|
|
|
|
def _stores(by_collection):
|
|
"""vectorstore(collection, cfg, pool) → a store whose
|
|
similarity_search_with_score_by_vector returns by_collection[name]."""
|
|
|
|
def factory(collection, cfg, pool):
|
|
s = MagicMock()
|
|
s.similarity_search_with_score_by_vector.return_value = by_collection.get(
|
|
collection, []
|
|
)
|
|
return s
|
|
|
|
return factory
|
|
|
|
|
|
class TestRetrieve:
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_merges_kinds_and_builds_links(self, mock_vs, MockEmb):
|
|
MockEmb.return_value.embed_query.return_value = [0.1] * 3
|
|
mock_vs.side_effect = _stores(
|
|
{
|
|
"comments": [
|
|
(
|
|
_doc(
|
|
"Telehealth comment.",
|
|
kind="comment",
|
|
comment_id="CMS-2026-2377-3438",
|
|
docket="CMS-2026-2377",
|
|
item_key="K1",
|
|
date="2026-08-19",
|
|
title="Anand M.",
|
|
),
|
|
0.25,
|
|
)
|
|
],
|
|
"rules": [
|
|
(
|
|
_doc(
|
|
"Under this proposal, the new G codes apply.",
|
|
kind="rule",
|
|
item_key="R1",
|
|
html_url="https://fr.test/doc",
|
|
p_id="935",
|
|
page="43949",
|
|
ordinal="1",
|
|
fr_volume="91",
|
|
date="2026-07-16",
|
|
title="CY2027 PFS NPRM",
|
|
),
|
|
0.20,
|
|
)
|
|
],
|
|
"corpus": [],
|
|
}
|
|
)
|
|
out = retrieve("telehealth", cfg=CFG, pool=MagicMock(), now=NOW)
|
|
assert [s["kind"] for s in out] == ["rule", "comment"]
|
|
rule, comment = out
|
|
assert rule["label"] == "91 FR 43949 ¶1"
|
|
assert rule["id"] == rule["label"]
|
|
assert rule["url"].startswith("https://fr.test/doc#p-935:~:text=Under%20this")
|
|
assert (
|
|
comment["url"] == "https://www.regulations.gov/comment/CMS-2026-2377-3438"
|
|
)
|
|
assert comment["comment_id"] == "CMS-2026-2377-3438"
|
|
assert comment["docket"] == "CMS-2026-2377"
|
|
assert comment["date"] == "2026-08-19"
|
|
assert comment["snippet"] == "Telehealth comment."
|
|
assert 0 < comment["score"] <= 1
|
|
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_over_fetches_three_per_kind_and_embeds_once(self, mock_vs, MockEmb):
|
|
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
|
|
retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
|
|
assert MockEmb.return_value.embed_query.call_count == 1
|
|
comments = stores["comments"].similarity_search_with_score_by_vector
|
|
rules = stores["rules"].similarity_search_with_score_by_vector
|
|
assert comments.call_args.kwargs["k"] == 6
|
|
assert rules.call_args.kwargs["k"] == 3
|
|
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_since_filters_old_hits(self, mock_vs, MockEmb):
|
|
MockEmb.return_value.embed_query.return_value = [0.0]
|
|
mock_vs.side_effect = _stores(
|
|
{
|
|
"comments": [
|
|
(
|
|
_doc(
|
|
"old",
|
|
kind="comment",
|
|
comment_id="C-1",
|
|
item_key="A",
|
|
date="2019-01-01",
|
|
),
|
|
0.1,
|
|
),
|
|
(
|
|
_doc(
|
|
"new",
|
|
kind="comment",
|
|
comment_id="C-2",
|
|
item_key="B",
|
|
date="2026-01-01",
|
|
),
|
|
0.3,
|
|
),
|
|
]
|
|
}
|
|
)
|
|
out = retrieve("q", cfg=CFG, pool=MagicMock(), since="2025-01-01", now=NOW)
|
|
assert [s["comment_id"] for s in out] == ["C-2"]
|
|
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_recent_comment_outranks_slightly_closer_old_one(self, mock_vs, MockEmb):
|
|
MockEmb.return_value.embed_query.return_value = [0.0]
|
|
mock_vs.side_effect = _stores(
|
|
{
|
|
"comments": [
|
|
(
|
|
_doc(
|
|
"old",
|
|
kind="comment",
|
|
comment_id="C-1",
|
|
item_key="A",
|
|
date="2019-01-01",
|
|
),
|
|
0.20,
|
|
),
|
|
(
|
|
_doc(
|
|
"new",
|
|
kind="comment",
|
|
comment_id="C-2",
|
|
item_key="B",
|
|
date="2026-08-19",
|
|
),
|
|
0.25,
|
|
),
|
|
]
|
|
}
|
|
)
|
|
out = retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
|
|
assert [s["comment_id"] for s in out] == ["C-2", "C-1"]
|
|
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_legacy_chunks_without_kind_are_treated_as_comments(self, mock_vs, MockEmb):
|
|
MockEmb.return_value.embed_query.return_value = [0.0]
|
|
mock_vs.side_effect = _stores(
|
|
{
|
|
"comments": [
|
|
(_doc("x", comment_id="C-9", item_key="K9", docket="D"), 0.5)
|
|
]
|
|
}
|
|
)
|
|
(s,) = retrieve("q", cfg=CFG, pool=MagicMock(), now=NOW)
|
|
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()
|
|
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_zero_k_for_a_kind_skips_its_collection_entirely(self, mock_vs, MockEmb):
|
|
# cfg.k_per_kind["corpus"] == 0 — _hits's "if k <= 0: continue"
|
|
# must skip that collection without ever opening its store.
|
|
MockEmb.return_value.embed_query.return_value = [0.0]
|
|
cfg0 = replace(CFG, k_per_kind={"comment": 0, "rule": 1, "corpus": 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
|
|
retrieve("q", cfg=cfg0, pool=MagicMock(), now=NOW)
|
|
assert set(stores) == {"rules"}
|
|
|
|
|
|
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 TestDocketRuleYear:
|
|
"""``_docket_rule_year`` degrades to era 0 for any unparsable
|
|
``comment_end_date`` rather than raising — the caller (``era_of``)
|
|
already treats 0 as "unknown, sorted last"."""
|
|
|
|
def test_unparsable_string_returns_zero(self):
|
|
assert rag._docket_rule_year("not-a-date") == 0
|
|
|
|
def test_none_returns_zero(self):
|
|
assert rag._docket_rule_year(None) == 0
|
|
|
|
|
|
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
|
|
|
|
def test_comment_empty_docket_short_circuits_without_store_lookup(self):
|
|
# "" is falsy — _docket_year's "if not docket_id: return 0" fires
|
|
# before ever touching _bib_store, so era_of falls straight back
|
|
# to the hit's own date year.
|
|
with patch("llm.rag._bib_store") as mock_store:
|
|
h = self._hit(kind="comment", docket="", date="2018-05-01")
|
|
assert era_of(h) == 2018
|
|
mock_store.assert_not_called()
|
|
|
|
@patch("llm.rag._bib_store", side_effect=RuntimeError("store boom"))
|
|
def test_bib_store_failure_degrades_to_own_date_year(self, mock_store, caplog):
|
|
# A broken bib.Store (open failure, bad mtime, bad query) must
|
|
# never break era-balanced retrieval over one docket lookup —
|
|
# it's logged and era_of falls back to the comment's own date.
|
|
h = self._hit(kind="comment", docket="CMS-2024-1", date="2024-03-01")
|
|
assert era_of(h) == 2024
|
|
assert "docket era lookup failed" in caplog.text
|
|
|
|
@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
|
|
|
|
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):
|
|
sources = [
|
|
{
|
|
"id": "91 FR 43949 ¶1",
|
|
"label": "91 FR 43949 ¶1",
|
|
"kind": "rule",
|
|
"date": "2026-07-16",
|
|
"snippet": "Under this proposal",
|
|
"url": "u",
|
|
"title": "t",
|
|
"docket": "",
|
|
"comment_id": "",
|
|
"score": 0.9,
|
|
},
|
|
{
|
|
"id": "CMS-2026-2377-1",
|
|
"label": "CMS-2026-2377-1",
|
|
"kind": "comment",
|
|
"date": "2026-08-19",
|
|
"snippet": "reduce documentation",
|
|
"url": "u",
|
|
"title": "t",
|
|
"docket": "CMS-2026-2377",
|
|
"comment_id": "CMS-2026-2377-1",
|
|
"score": 0.8,
|
|
},
|
|
]
|
|
msgs = build_messages("why?", sources)
|
|
sys_msg = msgs[0]["content"].lower()
|
|
assert msgs[0]["role"] == "system"
|
|
assert "only" in sys_msg
|
|
assert "don't have information" in sys_msg
|
|
assert "most recent" in sys_msg
|
|
assert (
|
|
"[91 FR 43949 ¶1] (rule, 2026-07-16) Under this proposal"
|
|
in msgs[1]["content"]
|
|
)
|
|
assert "[CMS-2026-2377-1] (comment, 2026-08-19)" in msgs[1]["content"]
|
|
assert "why?" in msgs[1]["content"]
|
|
|
|
def test_no_sources_marks_empty_context(self):
|
|
msgs = build_messages("q", [])
|
|
assert "no relevant excerpts" in msgs[1]["content"].lower()
|
|
|
|
def test_evidence_block_between_excerpts_and_question(self):
|
|
from llm.evidence import ValuationEvidence
|
|
from pfs.valuation import ValuationRow
|
|
|
|
row = ValuationRow(
|
|
code="G0556",
|
|
description="d",
|
|
vintage="CY2026 final",
|
|
year=2026,
|
|
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="[PFS CY2026 Addendum B]",
|
|
citation="90 FR 49266",
|
|
url="u",
|
|
)
|
|
ev = ValuationEvidence(("G0556",), ("APCM",), (row,), ())
|
|
msgs = build_messages("how much?", [], evidence=ev)
|
|
user = msgs[1]["content"]
|
|
assert (
|
|
user.index("Excerpts:")
|
|
< user.index("Valuation (authoritative")
|
|
< user.index("Question: how much?")
|
|
)
|
|
assert "[PFS CY2026 Addendum B] G0556" in user
|
|
assert "valuation" in msgs[0]["content"].lower()
|
|
assert "do not compute" in msgs[0]["content"].lower()
|
|
|
|
def test_no_evidence_prompt_unchanged(self):
|
|
assert build_messages("q", []) == build_messages("q", [], evidence=None)
|
|
assert "Valuation" not in build_messages("q", [])[1]["content"]
|
|
|
|
def test_lineage_block_between_excerpts_and_valuation(self):
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
le = LineageEvidence(
|
|
codes=("G2058",),
|
|
families=("CCM",),
|
|
events=(
|
|
LineageEvent(
|
|
"G2058",
|
|
2021,
|
|
"replaced_by",
|
|
(),
|
|
("99439",),
|
|
"CY2021 PFS final ¶686",
|
|
"YBM4IZUS",
|
|
686,
|
|
0,
|
|
"u",
|
|
"fr",
|
|
True,
|
|
"",
|
|
),
|
|
),
|
|
element_diffs=(),
|
|
guidance=(),
|
|
)
|
|
ev = ValuationEvidence(("G2058",), ("CCM",), (), ())
|
|
msgs = build_messages("history?", [], evidence=ev, lineage=le)
|
|
user = msgs[1]["content"]
|
|
assert (
|
|
user.index("Excerpts:")
|
|
< user.index("Lineage — dated facts")
|
|
< user.index("Valuation (authoritative")
|
|
< user.index("Question: history?")
|
|
)
|
|
assert "2021: G2058 was replaced by 99439 [CY2021 PFS final ¶686]" in user
|
|
assert "lineage" in msgs[0]["content"].lower()
|
|
assert "cite its bracketed label for every dated claim" in msgs[0]["content"]
|
|
|
|
def test_lineage_only_no_valuation_block(self):
|
|
le = LineageEvidence(
|
|
codes=("G2058",), families=(), events=(), element_diffs=(), guidance=()
|
|
)
|
|
msgs = build_messages("q", [], lineage=le)
|
|
user = msgs[1]["content"]
|
|
assert "Lineage — dated facts" in user
|
|
assert "Valuation (authoritative" not in user
|
|
|
|
def test_no_lineage_prompt_unchanged(self):
|
|
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_cited_and_manual_drop_after_earlier_blocks_hit_floor(self):
|
|
# Push the drop order all the way through: lineage-source (>4),
|
|
# retrieved (>6), cited (>8), manual (>1) — every source_steps
|
|
# reassignment branch, including "cited" and "manual".
|
|
retrieved = [_bmsrc(f"R{i}", snippet="R" * 20) for i in range(10)]
|
|
cited = [_bmsrc(f"C{i}", snippet="C" * 20) for i in range(10)]
|
|
ls = [_bmsrc(f"LS{i}", snippet="L" * 20) for i in range(8)]
|
|
manual = [_bmsrc(f"M{i}", snippet="M" * 20) for i in range(5)]
|
|
sources = retrieved + cited + ls + manual
|
|
target = build_messages(
|
|
"q",
|
|
retrieved[:6] + cited[:8] + ls[:4] + manual[:1],
|
|
n_retrieved=6,
|
|
n_cited=8,
|
|
n_lineage_sources=4,
|
|
n_manual=1,
|
|
)[1]["content"]
|
|
msgs = build_messages(
|
|
"q",
|
|
sources,
|
|
n_retrieved=10,
|
|
n_cited=10,
|
|
n_lineage_sources=8,
|
|
n_manual=5,
|
|
budget_chars=len(target),
|
|
)
|
|
assert msgs[1]["content"] == target
|
|
assert msgs[1]["content"].count("[C") == 8
|
|
assert msgs[1]["content"].count("[M") == 1
|
|
|
|
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()
|
|
pool.acquire_generation.return_value.__enter__.return_value = "http://h1:11434"
|
|
pool.vram.return_value = vram
|
|
pool.serves.return_value = serves_big
|
|
return pool
|
|
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.valuation_evidence", return_value=None)
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_yields_tokens_then_sources_then_done(
|
|
self, mock_retrieve, MockClient, _lin, _ev, mock_engine
|
|
):
|
|
src = {
|
|
"id": "C1",
|
|
"label": "C1",
|
|
"kind": "comment",
|
|
"snippet": "s",
|
|
"score": 0.1,
|
|
}
|
|
mock_retrieve.return_value = [src]
|
|
lines = [
|
|
'{"message":{"content":"Doc"},"done":false}',
|
|
"", # keep-alive blank line — must be skipped, not parsed
|
|
'{"message":{"content":"tors"},"done":false}',
|
|
'{"message":{"content":""},"done":true}',
|
|
]
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(lines)
|
|
pool = self._pool()
|
|
|
|
events = list(stream_answer("q", cfg=CFG, pool=pool))
|
|
|
|
pool.check.assert_called_once_with("chat")
|
|
# no codes in the question — pgvector is never asked for cited rules
|
|
mock_engine.assert_not_called()
|
|
mock_retrieve.assert_called_once_with(
|
|
"q", cfg=CFG, pool=pool, since="", mode="auto", exclude_phrases=()
|
|
)
|
|
assert events[0] == {"type": "token", "text": "Doc"}
|
|
assert events[1] == {"type": "token", "text": "tors"}
|
|
assert events[-2] == {
|
|
"type": "sources",
|
|
"sources": [src],
|
|
"model": "big",
|
|
"host": "http://h1:11434",
|
|
"mode": "recent",
|
|
}
|
|
assert events[-1] == {"type": "done"}
|
|
body = client.stream.call_args.kwargs["json"]
|
|
assert body["model"] == "big"
|
|
assert body["options"] == {
|
|
"num_ctx": 8192,
|
|
"temperature": CFG.chat_temperature,
|
|
"num_predict": CFG.chat_max_tokens,
|
|
}
|
|
assert client.stream.call_args.args[1] == "http://h1:11434/api/chat"
|
|
|
|
@patch("llm.rag.valuation_evidence", return_value=None)
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_small_host_uses_baseline_model(self, mock_retrieve, MockClient, _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(vram=12.0)))
|
|
assert client.stream.call_args.kwargs["json"]["model"] == "chat"
|
|
assert events[-2]["model"] == "chat"
|
|
|
|
@patch("llm.rag.valuation_evidence", return_value=None)
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_since_forwarded(self, mock_retrieve, MockClient, _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}'])
|
|
list(stream_answer("q", cfg=CFG, pool=self._pool(), 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.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_http_error_propagates(self, mock_retrieve, MockClient, _lin, _ev):
|
|
mock_retrieve.return_value = []
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.raise_for_status.side_effect = RuntimeError("ollama down")
|
|
with pytest.raises(RuntimeError, match="ollama down"):
|
|
list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
|
|
|
@patch("llm.rag._manual_sources", return_value=[])
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_valuation_event_first_and_sources_merged(
|
|
self,
|
|
mock_retrieve,
|
|
MockClient,
|
|
_lin,
|
|
mock_ev,
|
|
mock_cited,
|
|
mock_engine,
|
|
_manual,
|
|
):
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
src = {
|
|
"id": "C1",
|
|
"label": "C1",
|
|
"kind": "comment",
|
|
"snippet": "s",
|
|
"score": 0.1,
|
|
}
|
|
cited = {
|
|
"id": "89 FR 97710 ¶3",
|
|
"label": "89 FR 97710 ¶3",
|
|
"kind": "rule",
|
|
"snippet": "G0556",
|
|
"score": 0.0,
|
|
}
|
|
mock_retrieve.return_value = [src]
|
|
mock_cited.return_value = [cited]
|
|
mock_ev.return_value = ValuationEvidence(("G0556",), ("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}'])
|
|
|
|
events = list(stream_answer("APCM?", cfg=CFG, pool=self._pool()))
|
|
|
|
assert events[0]["type"] == "valuation" and events[0]["codes"] == ["G0556"]
|
|
assert events[1] == {"type": "token", "text": "x"}
|
|
assert events[-2]["sources"] == [src, cited]
|
|
mock_cited.assert_called_once_with(
|
|
mock_engine.return_value,
|
|
("G0556",),
|
|
per_code=CFG.code_cited_per_code,
|
|
collections=CFG.code_cited_collections,
|
|
families=("APCM",),
|
|
max_total=CFG.code_cited_max,
|
|
by_docket=False,
|
|
)
|
|
body = client.stream.call_args.kwargs["json"]
|
|
assert "Valuation (authoritative" in body["messages"][1]["content"]
|
|
|
|
@patch("llm.rag._manual_sources", return_value=[])
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_conflicting_cited_snippet_is_dropped_and_logged(
|
|
self,
|
|
mock_retrieve,
|
|
MockClient,
|
|
_lin,
|
|
mock_ev,
|
|
mock_cited,
|
|
mock_engine,
|
|
_manual,
|
|
caplog,
|
|
):
|
|
# A code-cited excerpt (not filtered by retrieve()'s own
|
|
# exclude_phrases, since it never goes through retrieve) that
|
|
# uses "CCM" for cardiac contractility modulation must still be
|
|
# dropped by the final _drop_conflicts pass, and the drop logged.
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
clean = {
|
|
"id": "C1",
|
|
"label": "C1",
|
|
"kind": "rule",
|
|
"snippet": "CCM services require a comprehensive care plan",
|
|
"score": 0.1,
|
|
}
|
|
conflicting = {
|
|
"id": "C2",
|
|
"label": "C2",
|
|
"kind": "rule",
|
|
"snippet": "Cardiac Contractility Modulation (CCM) pulse generator",
|
|
"score": 0.2,
|
|
}
|
|
mock_retrieve.return_value = []
|
|
mock_cited.return_value = [clean, conflicting]
|
|
mock_ev.return_value = ValuationEvidence((), ("CCM",), (), ())
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
|
|
|
|
with caplog.at_level(logging.INFO, logger="llm.rag"):
|
|
events = list(stream_answer("CCM?", cfg=CFG, pool=self._pool()))
|
|
|
|
assert events[-2]["sources"] == [clean]
|
|
assert "dropped 1 excerpt(s) using a family acronym" in caplog.text
|
|
|
|
@patch("llm.rag._manual_sources", return_value=[])
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.lineage_evidence")
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_lineage_event_before_valuation_and_tokens(
|
|
self,
|
|
mock_retrieve,
|
|
MockClient,
|
|
mock_lin,
|
|
mock_ev,
|
|
mock_cited,
|
|
mock_engine,
|
|
_manual,
|
|
):
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
mock_retrieve.return_value = []
|
|
mock_cited.return_value = []
|
|
mock_lin.return_value = LineageEvidence(
|
|
codes=("G2058",),
|
|
families=("CCM",),
|
|
events=(
|
|
LineageEvent(
|
|
"G2058",
|
|
2021,
|
|
"replaced_by",
|
|
(),
|
|
("99439",),
|
|
"CY2021 PFS final ¶686",
|
|
"YBM4IZUS",
|
|
686,
|
|
0,
|
|
"u",
|
|
"fr",
|
|
True,
|
|
"",
|
|
),
|
|
),
|
|
element_diffs=(),
|
|
guidance=(),
|
|
)
|
|
mock_ev.return_value = ValuationEvidence(("G2058",), ("CCM",), (), ())
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
|
|
|
|
events = list(stream_answer("history of G2058?", cfg=CFG, pool=self._pool()))
|
|
|
|
assert [e["type"] for e in events] == [
|
|
"lineage",
|
|
"valuation",
|
|
"token",
|
|
"sources",
|
|
"done",
|
|
]
|
|
assert events[0]["codes"] == ["G2058"]
|
|
body = client.stream.call_args.kwargs["json"]
|
|
content = body["messages"][1]["content"]
|
|
assert content.index("Lineage — dated facts") < content.index(
|
|
"Valuation (authoritative"
|
|
)
|
|
|
|
@patch("llm.rag._manual_sources", return_value=[])
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence", return_value=None)
|
|
@patch("llm.rag.lineage_evidence")
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_lineage_only_calls_cited_sources_with_lineage_codes(
|
|
self, mock_retrieve, MockClient, mock_lin, _ev, mock_cited, mock_engine, _manual
|
|
):
|
|
mock_retrieve.return_value = []
|
|
mock_cited.return_value = []
|
|
mock_lin.return_value = LineageEvidence(
|
|
codes=("99490", "99491"),
|
|
families=("CCM",),
|
|
events=(),
|
|
element_diffs=(),
|
|
guidance=(),
|
|
)
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
|
|
|
|
events = list(stream_answer("CCM history?", cfg=CFG, pool=self._pool()))
|
|
|
|
assert events[0]["type"] == "lineage"
|
|
# "CCM history?" is history-shaped (is_history_question) — mode
|
|
# resolves to "timeline", so the comments window is per-docket.
|
|
mock_cited.assert_called_once_with(
|
|
mock_engine.return_value,
|
|
("99490", "99491"),
|
|
per_code=CFG.code_cited_per_code,
|
|
collections=CFG.code_cited_collections,
|
|
families=("CCM",),
|
|
max_total=CFG.code_cited_max,
|
|
by_docket=True,
|
|
)
|
|
|
|
@patch("llm.rag._manual_sources", return_value=[])
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.lineage_evidence")
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_codes_and_families_unioned_when_both_present(
|
|
self,
|
|
mock_retrieve,
|
|
MockClient,
|
|
mock_lin,
|
|
mock_ev,
|
|
mock_cited,
|
|
mock_engine,
|
|
_manual,
|
|
):
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
mock_retrieve.return_value = []
|
|
mock_cited.return_value = []
|
|
# lineage reaches a code (G2058) with no RVU rows valuation never sees.
|
|
mock_lin.return_value = LineageEvidence(
|
|
codes=("99490", "G2058"),
|
|
families=("CCM",),
|
|
events=(),
|
|
element_diffs=(),
|
|
guidance=(),
|
|
)
|
|
mock_ev.return_value = ValuationEvidence(("99490",), ("CCM",), (), ())
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
|
|
|
|
list(stream_answer("CCM and G2058?", cfg=CFG, pool=self._pool()))
|
|
|
|
mock_cited.assert_called_once_with(
|
|
mock_engine.return_value,
|
|
("99490", "G2058"),
|
|
per_code=CFG.code_cited_per_code,
|
|
collections=CFG.code_cited_collections,
|
|
families=("CCM",),
|
|
max_total=CFG.code_cited_max,
|
|
by_docket=False,
|
|
)
|
|
|
|
@patch("llm.rag._manual_sources")
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_manual_sources_merged_after_cited_sources(
|
|
self,
|
|
mock_retrieve,
|
|
MockClient,
|
|
_lin,
|
|
mock_ev,
|
|
mock_cited,
|
|
mock_engine,
|
|
mock_manual,
|
|
):
|
|
"""#691/#705: the CPT manual's guideline source is a third,
|
|
deterministic layer — merged after the retrieved and code-cited
|
|
sources, never before them."""
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
src = {
|
|
"id": "C1",
|
|
"label": "C1",
|
|
"kind": "comment",
|
|
"snippet": "s",
|
|
"score": 0.1,
|
|
}
|
|
cited = {
|
|
"id": "89 FR 97710 ¶3",
|
|
"label": "89 FR 97710 ¶3",
|
|
"kind": "rule",
|
|
"snippet": "G0556",
|
|
"score": 0.0,
|
|
}
|
|
manual = {
|
|
"id": "CPT 2024 — Advanced Primary Care Management",
|
|
"label": "CPT 2024 — Advanced Primary Care Management",
|
|
"kind": "corpus",
|
|
"snippet": "guideline text",
|
|
"score": 0.0,
|
|
}
|
|
mock_retrieve.return_value = [src]
|
|
mock_cited.return_value = [cited]
|
|
mock_manual.return_value = [manual]
|
|
mock_ev.return_value = ValuationEvidence(("G0556",), ("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}'])
|
|
|
|
events = list(stream_answer("APCM?", cfg=CFG, pool=self._pool()))
|
|
|
|
assert events[-2]["sources"] == [src, cited, manual]
|
|
mock_manual.assert_called_once_with(("APCM",), CFG)
|
|
|
|
@patch("llm.rag._manual_sources", return_value=[])
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.lineage_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_manual_sources_skips_wide_families(
|
|
self,
|
|
mock_retrieve,
|
|
MockClient,
|
|
_lin,
|
|
mock_ev,
|
|
mock_cited,
|
|
mock_engine,
|
|
mock_manual,
|
|
):
|
|
from llm.evidence import ValuationEvidence
|
|
|
|
mock_retrieve.return_value = []
|
|
mock_cited.return_value = []
|
|
mock_ev.return_value = ValuationEvidence(
|
|
(), ("BIGFAM", "APCM"), (), (), wide=("BIGFAM",)
|
|
)
|
|
client = MockClient.return_value.__enter__.return_value
|
|
resp = client.stream.return_value.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(['{"message":{"content":"x"},"done":true}'])
|
|
|
|
list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
|
|
|
mock_manual.assert_called_once_with(("APCM",), CFG)
|
|
|
|
@patch("llm.rag._manual_sources")
|
|
@patch("llm.rag._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@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_no_evidence_no_lineage_never_calls_manual_sources(
|
|
self, mock_retrieve, MockClient, _lin, _ev, mock_cited, mock_engine, mock_manual
|
|
):
|
|
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":"x"},"done":true}'])
|
|
|
|
list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
|
|
|
mock_cited.assert_not_called()
|
|
mock_manual.assert_not_called()
|
|
|
|
@patch("llm.rag.valuation_evidence", return_value=None)
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_control_question_events_match_the_no_lineage_baseline(
|
|
self, mock_retrieve, MockClient, _ev
|
|
):
|
|
"""No codes in the question — the real lineage_evidence short-
|
|
circuits on an empty Detection without touching the replica, so
|
|
its event stream must be byte-identical to one where the feature
|
|
is switched off outright (``lineage_evidence`` patched to
|
|
``None``). 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
|
|
|
|
def _run():
|
|
resp.iter_lines.return_value = iter(
|
|
['{"message":{"content":"x"},"done":true}']
|
|
)
|
|
return list(
|
|
stream_answer(
|
|
"why did CMS finalize this policy?", cfg=cfg, pool=self._pool()
|
|
)
|
|
)
|
|
|
|
with_real_lineage = json.dumps(_run())
|
|
with patch("llm.rag.lineage_evidence", return_value=None):
|
|
with_lineage_off = json.dumps(_run())
|
|
assert with_real_lineage == with_lineage_off
|
|
|
|
|
|
class TestManualSourcesWrapper:
|
|
"""``_manual_sources`` — the cursor-per-call wrapper around
|
|
``evidence.manual_sources`` (mirrors ``valuation_evidence``'s own
|
|
cached-replica-cursor pattern)."""
|
|
|
|
def test_empty_families_short_circuits_without_touching_the_replica(self):
|
|
with patch("llm.rag._connect") as mock_connect:
|
|
assert rag._manual_sources((), CFG) == []
|
|
mock_connect.assert_not_called()
|
|
|
|
@patch("llm.rag.manual_sources")
|
|
@patch("llm.rag._connect")
|
|
@patch("llm.rag._replica_path", return_value="/x/aco.ro.duckdb")
|
|
def test_opens_a_cursor_calls_through_and_closes_it(
|
|
self, _path, mock_connect, mock_manual
|
|
):
|
|
cur = MagicMock()
|
|
mock_connect.return_value.cursor.return_value = cur
|
|
mock_manual.return_value = [{"label": "CPT 2024 — X"}]
|
|
|
|
out = rag._manual_sources(("CCM",), CFG)
|
|
|
|
assert out == [{"label": "CPT 2024 — X"}]
|
|
mock_connect.assert_called_once_with("/x/aco.ro.duckdb")
|
|
mock_manual.assert_called_once_with(cur, ("CCM",))
|
|
cur.close.assert_called_once() # the cached parent connection stays open
|
|
|
|
@patch("llm.rag._connect", side_effect=OSError("no replica"))
|
|
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()
|
|
|
|
|
|
class TestHasConflict:
|
|
def test_no_phrases_returns_false(self):
|
|
from llm.rag import _has_conflict
|
|
|
|
assert _has_conflict("some text mentioning ccm", ()) is False
|
|
|
|
def test_empty_text_returns_false(self):
|
|
from llm.rag import _has_conflict
|
|
|
|
assert _has_conflict("", ("cardiac contractility modulation",)) is False
|
|
|
|
|
|
class TestAcronymConflicts:
|
|
"""CCM = chronic care management to the family registry, cardiac
|
|
contractility modulation to the CY2027 NPRM's "CCM code family";
|
|
excerpts using the acronym the other way are dropped."""
|
|
|
|
def test_phrases_come_from_the_family_registry(self):
|
|
from llm.rag import _conflict_phrases
|
|
|
|
phrases = _conflict_phrases(["CCM"])
|
|
assert "cardiac contractility modulation" in phrases
|
|
assert _conflict_phrases(["TCM", "NOPE"]) == ()
|
|
|
|
def test_drop_conflicts_on_snippets(self):
|
|
from llm.rag import _drop_conflicts
|
|
|
|
srcs = [
|
|
{
|
|
"label": "a",
|
|
"snippet": "the CCM code family for Cardiac Contractility Modulation",
|
|
},
|
|
{"label": "b", "snippet": "CCM services require a comprehensive care plan"},
|
|
]
|
|
kept, n = _drop_conflicts(srcs, ("cardiac contractility modulation",))
|
|
assert [s["label"] for s in kept] == ["b"] and n == 1
|
|
assert _drop_conflicts(srcs, ()) == (srcs, 0)
|
|
|
|
@patch("llm.rag.PoolEmbeddings")
|
|
@patch("llm.index.vectorstore")
|
|
def test_retrieve_excludes_conflicting_hits_before_ranking(self, mock_vs, mock_emb):
|
|
from llm.rag import retrieve
|
|
|
|
mock_emb.return_value.embed_query.return_value = [0.1]
|
|
hits = {
|
|
"rules": [
|
|
(
|
|
_doc(
|
|
"Cardiac contractility modulation (CCM) pulse generator",
|
|
item_key="X1",
|
|
),
|
|
0.1,
|
|
),
|
|
(_doc("CCM services: comprehensive care plan", item_key="X2"), 0.2),
|
|
]
|
|
}
|
|
mock_vs.side_effect = _stores(hits)
|
|
out = retrieve(
|
|
"q",
|
|
cfg=CFG,
|
|
pool=MagicMock(),
|
|
mode="recent",
|
|
exclude_phrases=("cardiac contractility modulation",),
|
|
)
|
|
assert [s["label"] for s in out] and all(
|
|
"contractility" not in s["snippet"].lower() for s in out
|
|
)
|