1077 lines
39 KiB
Python
1077 lines
39 KiB
Python
"""llm.rag — multi-collection retrieval + grounded streaming answer."""
|
|
|
|
import json
|
|
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()
|
|
|
|
|
|
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:
|
|
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 events")
|
|
< user.index("Valuation (authoritative")
|
|
< user.index("Question: history?")
|
|
)
|
|
assert "[CY2021 PFS final ¶686] 2021 replaced_by G2058" in user
|
|
assert "lineage" in msgs[0]["content"].lower()
|
|
assert "cite its bracketed label for every dated claim" in msgs[0]["content"]
|
|
|
|
def test_lineage_only_no_valuation_block(self):
|
|
le = LineageEvidence(
|
|
codes=("G2058",), families=(), events=(), element_diffs=(), guidance=()
|
|
)
|
|
msgs = build_messages("q", [], lineage=le)
|
|
user = msgs[1]["content"]
|
|
assert "Lineage (dated events" in user
|
|
assert "Valuation (authoritative" not in user
|
|
|
|
def test_no_lineage_prompt_unchanged(self):
|
|
assert build_messages("q", []) == build_messages("q", [], lineage=None)
|
|
|
|
|
|
class TestStreamAnswer:
|
|
def _pool(self, vram=24.0, serves_big=True):
|
|
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"
|
|
)
|
|
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}
|
|
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,
|
|
)
|
|
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")
|
|
@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 events") < 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"
|
|
mock_cited.assert_called_once_with(
|
|
mock_engine.return_value,
|
|
("99490", "99491"),
|
|
per_code=CFG.code_cited_per_code,
|
|
collections=CFG.code_cited_collections,
|
|
families=("CCM",),
|
|
max_total=CFG.code_cited_max,
|
|
)
|
|
|
|
@patch("llm.rag._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,
|
|
)
|
|
|
|
@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.detect_codes")
|
|
@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,
|
|
mock_detect,
|
|
):
|
|
from llm.evidence import ValuationEvidence
|
|
from pfs.families import Detection
|
|
|
|
mock_retrieve.return_value = []
|
|
mock_cited.return_value = []
|
|
mock_detect.return_value = Detection(
|
|
codes=(), families=("BIGFAM", "APCM"), explicit=(), wide=("BIGFAM",)
|
|
)
|
|
mock_ev.return_value = ValuationEvidence((), ("BIGFAM", "APCM"), (), ())
|
|
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``)."""
|
|
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
|