31d9f27d11
code_cited_sources gains max_total (default 12, <=0 = unlimited): after the existing per-collection caps and dedupe, code-cited rows are round-robin interleaved across the configured collections in order, then family-cited rows the same way, appended after — so one chatty collection can no longer crowd the others out ahead of a truncation cut. New [llm] knob code_cited_max wires it from config into stream_answer. Also guards code_cited_collections against a bare scalar string in stack.toml collapsing into a tuple of characters.
425 lines
16 KiB
Python
425 lines
16 KiB
Python
"""llm.rag — multi-collection retrieval + grounded streaming answer."""
|
|
|
|
from datetime import date
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from langchain_core.documents import Document
|
|
|
|
from llm.config import LlmConfig
|
|
from llm.rag import build_messages, retrieve, stream_answer
|
|
|
|
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)
|
|
|
|
|
|
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"
|
|
|
|
|
|
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"]
|
|
|
|
|
|
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.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_yields_tokens_then_sources_then_done(
|
|
self, mock_retrieve, MockClient, _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="")
|
|
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",
|
|
}
|
|
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.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_small_host_uses_baseline_model(self, mock_retrieve, MockClient, _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.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_since_forwarded(self, mock_retrieve, MockClient, _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.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_http_error_propagates(self, mock_retrieve, MockClient, _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._engine")
|
|
@patch("llm.rag.code_cited_sources")
|
|
@patch("llm.rag.valuation_evidence")
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_valuation_event_first_and_sources_merged(
|
|
self, mock_retrieve, MockClient, mock_ev, mock_cited, mock_engine
|
|
):
|
|
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"]
|