52ce69acf1
Some checks failed
CI / notebooks-smoke (push) Successful in 1m34s
CI / test (push) Failing after 2m19s
Infra CI / zotero (push) Successful in 16s
Infra CI / docs (push) Successful in 17s
Infra CI / api (push) Successful in 1m7s
CI / lint (push) Successful in 38s
Infra CI / llm (push) Successful in 1m10s
Infra CI / mc (push) Failing after 41s
CI / test (pull_request) Failing after 2m27s
Infra CI / notebooks (push) Successful in 7m51s
CI / notebooks-smoke (pull_request) Successful in 1m25s
Infra CI / zotero (pull_request) Successful in 14s
Infra CI / docs (pull_request) Successful in 19s
Infra CI / notebooks (pull_request) Successful in 51s
CI / lint (pull_request) Successful in 32s
Infra CI / api (pull_request) Successful in 21s
Infra CI / llm (pull_request) Successful in 17s
Infra CI / mc (pull_request) Failing after 1m8s
The excerpt budget was a flat chat_num_ctx * 3 characters on the user
message alone. On FR/regulatory text the instruct model's tokenizer yields
2.5–2.9 chars per token, and the system prompt and the 1,200-token reply
reservation were not counted, so a routine answer rendered 9.4k–10k prompt
tokens into an 8,192-token context. Ollama's context-shift truncation keeps
the first 4 tokens and the tail half: the system prompt and the best-ranked
excerpts were dropped on every such answer, the only trace a WARN in the
ollama log. The nightly golden set has reported missing anchors since.
- prompt_budget_chars(cfg): (chat_num_ctx - chat_max_tokens - template
reserve) * chat_chars_per_token - len(system prompt); new [llm] knob
chat_chars_per_token = 2.4 (measured minimum 2.50 on the golden set)
- /api/chat is sent with truncate: false; the host's 400 for an over-long
prompt (exceed_context_size_error with n_prompt_tokens / n_ctx, or the
older bare message) rebuilds the messages under a proportionally tighter
budget and resends; when the trimming floors cannot shorten the prompt,
or on the final attempt, host truncation is allowed as the last resort
- stack_llm_chat_prompt_overflow_total{resolution} counts rebuilt /
host_truncated; each overflow is logged with the host's token counts
1705 lines
64 KiB
Python
1705 lines
64 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 (
|
||
_SYSTEM,
|
||
build_messages,
|
||
era_balance,
|
||
era_of,
|
||
is_history_question,
|
||
prompt_budget_chars,
|
||
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
|
||
)
|
||
|
||
|
||
class TestPromptBudget:
|
||
"""The excerpt budget is derived from the token context, not a flat
|
||
multiple of it: the reply reservation, the chat template and the
|
||
system prompt all come out first, so system + user + reply fit
|
||
``chat_num_ctx`` and the host never has to cut the prompt."""
|
||
|
||
def test_budget_leaves_room_for_system_prompt_and_reply(self):
|
||
from llm import rag
|
||
|
||
tokens = (
|
||
CFG.chat_num_ctx - CFG.chat_max_tokens - rag._PROMPT_TEMPLATE_RESERVE_TOKENS
|
||
)
|
||
assert prompt_budget_chars(CFG) == int(tokens * CFG.chat_chars_per_token) - len(
|
||
_SYSTEM
|
||
)
|
||
# strictly below the old flat 3× multiple that overflowed an 8k context
|
||
assert prompt_budget_chars(CFG) < CFG.chat_num_ctx * 3
|
||
|
||
def test_budget_scales_with_context_and_ratio(self):
|
||
from dataclasses import replace
|
||
|
||
wider = replace(CFG, chat_num_ctx=CFG.chat_num_ctx * 2)
|
||
assert prompt_budget_chars(wider) > prompt_budget_chars(CFG)
|
||
denser = replace(CFG, chat_chars_per_token=1.0)
|
||
assert prompt_budget_chars(denser) < prompt_budget_chars(CFG)
|
||
|
||
def test_budget_never_negative(self):
|
||
from dataclasses import replace
|
||
|
||
tiny = replace(CFG, chat_num_ctx=16, chat_max_tokens=8)
|
||
assert prompt_budget_chars(tiny) == 0
|
||
|
||
|
||
def _overflow_resp(prompt_tokens: int | None, num_ctx: int = 8192):
|
||
"""A mocked ``/api/chat`` 400 as Ollama returns it with ``truncate:
|
||
false`` — the current nested form with token counts, or (None) the
|
||
older bare message."""
|
||
resp = MagicMock()
|
||
resp.status_code = 400
|
||
if prompt_tokens is None:
|
||
body = {
|
||
"error": "the prompt is longer than the context length currently available to the model"
|
||
}
|
||
else:
|
||
inner = {
|
||
"error": {
|
||
"code": 400,
|
||
"message": f"request ({prompt_tokens} tokens) exceeds the available context size ({num_ctx} tokens), try increasing it",
|
||
"type": "exceed_context_size_error",
|
||
"n_prompt_tokens": prompt_tokens,
|
||
"n_ctx": num_ctx,
|
||
}
|
||
}
|
||
body = {"error": json.dumps(inner)}
|
||
resp.read.return_value = json.dumps(body).encode()
|
||
return resp
|
||
|
||
|
||
def _ok_resp(
|
||
lines=(
|
||
'{"message":{"content":"ok"},"done":false}',
|
||
'{"message":{"content":""},"done":true,"prompt_eval_count":5000,"eval_count":1}',
|
||
),
|
||
):
|
||
resp = MagicMock()
|
||
resp.status_code = 200
|
||
resp.iter_lines.return_value = iter(lines)
|
||
return resp
|
||
|
||
|
||
class TestStreamAnswerPromptOverflow:
|
||
"""The chat sends ``truncate: false`` and sizes the prompt itself; a
|
||
host 400 for an over-long prompt rebuilds the messages under a
|
||
tighter budget (from the host's own token counts when it reports
|
||
them) and resends — host-side truncation is the last resort only."""
|
||
|
||
def _pool(self):
|
||
pool = MagicMock()
|
||
pool.acquire_generation.return_value.__enter__.return_value = "http://h1:11434"
|
||
pool.vram.return_value = 24.0
|
||
pool.serves.return_value = True
|
||
return pool
|
||
|
||
def _src(self, i):
|
||
return {
|
||
"id": f"C{i}",
|
||
"label": f"C{i}",
|
||
"kind": "comment",
|
||
"snippet": "x" * 400,
|
||
"score": 0.1,
|
||
}
|
||
|
||
@patch("llm.rag.metrics.chat_prompt_overflow")
|
||
@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_first_attempt_uses_token_budget_and_refuses_host_truncation(
|
||
self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
|
||
):
|
||
mock_retrieve.return_value = [self._src(0)]
|
||
client = MockClient.return_value.__enter__.return_value
|
||
client.stream.return_value.__enter__.return_value = _ok_resp()
|
||
|
||
with patch("llm.rag.build_messages", wraps=build_messages) as bm:
|
||
events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
||
|
||
assert events[0] == {"type": "token", "text": "ok"}
|
||
assert bm.call_args.kwargs["budget_chars"] == prompt_budget_chars(CFG)
|
||
body = client.stream.call_args.kwargs["json"]
|
||
assert body["truncate"] is False
|
||
assert client.stream.call_count == 1
|
||
mock_overflow.assert_not_called()
|
||
|
||
@patch("llm.rag.metrics.chat_prompt_overflow")
|
||
@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_overflow_rebuilds_tighter_from_host_counts_then_streams(
|
||
self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
|
||
):
|
||
from llm import rag
|
||
|
||
mock_retrieve.return_value = [self._src(i) for i in range(30)]
|
||
client = MockClient.return_value.__enter__.return_value
|
||
client.stream.return_value.__enter__.side_effect = [
|
||
_overflow_resp(prompt_tokens=9500, num_ctx=8192),
|
||
_ok_resp(),
|
||
]
|
||
|
||
with patch("llm.rag.build_messages", wraps=build_messages) as bm:
|
||
events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
||
|
||
assert events[0] == {"type": "token", "text": "ok"}
|
||
assert events[-1] == {"type": "done"}
|
||
assert client.stream.call_count == 2
|
||
first, second = (c.kwargs["budget_chars"] for c in bm.call_args_list)
|
||
fit = 8192 - CFG.chat_max_tokens - rag._PROMPT_TEMPLATE_RESERVE_TOKENS
|
||
assert first == prompt_budget_chars(CFG)
|
||
assert second == int(first * fit / 9500 * rag._OVERFLOW_MARGIN)
|
||
assert second < first
|
||
# the second request still refuses host-side truncation
|
||
assert client.stream.call_args_list[1].kwargs["json"]["truncate"] is False
|
||
# the resent prompt really is smaller
|
||
sent = [
|
||
c.kwargs["json"]["messages"][1]["content"]
|
||
for c in client.stream.call_args_list
|
||
]
|
||
assert len(sent[1]) < len(sent[0])
|
||
mock_overflow.assert_called_once_with("rebuilt")
|
||
|
||
@patch("llm.rag.metrics.chat_prompt_overflow")
|
||
@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_bare_message_overflow_shrinks_by_fixed_factor(
|
||
self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
|
||
):
|
||
from llm import rag
|
||
|
||
mock_retrieve.return_value = [self._src(i) for i in range(30)]
|
||
client = MockClient.return_value.__enter__.return_value
|
||
client.stream.return_value.__enter__.side_effect = [
|
||
_overflow_resp(prompt_tokens=None),
|
||
_ok_resp(),
|
||
]
|
||
|
||
with patch("llm.rag.build_messages", wraps=build_messages) as bm:
|
||
list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
||
|
||
first, second = (c.kwargs["budget_chars"] for c in bm.call_args_list)
|
||
assert second == int(first * rag._OVERFLOW_SHRINK)
|
||
mock_overflow.assert_called_once_with("rebuilt")
|
||
|
||
@patch("llm.rag.metrics.chat_prompt_overflow")
|
||
@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_repeated_overflow_falls_back_to_host_truncation(
|
||
self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
|
||
):
|
||
mock_retrieve.return_value = [self._src(i) for i in range(30)]
|
||
client = MockClient.return_value.__enter__.return_value
|
||
client.stream.return_value.__enter__.side_effect = [
|
||
_overflow_resp(9500),
|
||
_overflow_resp(9000),
|
||
_ok_resp(),
|
||
]
|
||
|
||
events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
||
|
||
assert events[-1] == {"type": "done"}
|
||
flags = [c.kwargs["json"]["truncate"] for c in client.stream.call_args_list]
|
||
assert flags == [False, False, True]
|
||
assert [c.args[0] for c in mock_overflow.call_args_list] == [
|
||
"rebuilt",
|
||
"host_truncated",
|
||
]
|
||
|
||
@patch("llm.rag.metrics.chat_prompt_overflow")
|
||
@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_overflow_at_the_trimming_floors_goes_straight_to_host_truncation(
|
||
self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
|
||
):
|
||
# One excerpt: the floors keep it however small the budget gets, so a
|
||
# second truncate:false attempt would only bounce again.
|
||
mock_retrieve.return_value = [self._src(0)]
|
||
client = MockClient.return_value.__enter__.return_value
|
||
client.stream.return_value.__enter__.side_effect = [
|
||
_overflow_resp(9500),
|
||
_ok_resp(),
|
||
]
|
||
|
||
events = list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
||
|
||
assert events[-1] == {"type": "done"}
|
||
flags = [c.kwargs["json"]["truncate"] for c in client.stream.call_args_list]
|
||
assert flags == [False, True]
|
||
mock_overflow.assert_called_once_with("host_truncated")
|
||
|
||
@patch("llm.rag.metrics.chat_prompt_overflow")
|
||
@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_unrelated_400_still_raises(
|
||
self, mock_retrieve, MockClient, _lin, _ev, mock_overflow
|
||
):
|
||
import httpx
|
||
|
||
mock_retrieve.return_value = [self._src(0)]
|
||
resp = MagicMock()
|
||
resp.status_code = 400
|
||
resp.read.return_value = json.dumps({"error": "model 'big' not found"}).encode()
|
||
resp.raise_for_status.side_effect = httpx.HTTPStatusError(
|
||
"400", request=MagicMock(), response=MagicMock()
|
||
)
|
||
client = MockClient.return_value.__enter__.return_value
|
||
client.stream.return_value.__enter__.return_value = resp
|
||
|
||
with pytest.raises(httpx.HTTPStatusError):
|
||
list(stream_answer("q", cfg=CFG, pool=self._pool()))
|
||
assert client.stream.call_count == 1
|
||
mock_overflow.assert_not_called()
|