edb72e9aba
New llm FastAPI service (src/llm/api.py + rag.py + web/chat.html): grounded streaming chat over indexed comments with cited sources. Own image, compose service, Traefik reef entry with git-sso, llm subdomain registered. Dashboard tile + README row. stack llm serve CLI.
126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
"""llm.rag — retrieval + grounded streaming answer."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
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",),
|
|
embed_model="embed",
|
|
instruct_model="chat",
|
|
embed_dim=768,
|
|
pg_host="x",
|
|
pg_port=5432,
|
|
pg_db="llm",
|
|
pg_user="llm",
|
|
build_ann_index=True,
|
|
)
|
|
|
|
|
|
def _doc(text, **md):
|
|
return Document(page_content=text, metadata=md)
|
|
|
|
|
|
class TestRetrieve:
|
|
@patch("llm.index.vectorstore")
|
|
def test_maps_hits_to_sources(self, mock_vs):
|
|
store = mock_vs.return_value
|
|
store.similarity_search_with_score.return_value = [
|
|
(
|
|
_doc(
|
|
"Telehealth comment.",
|
|
comment_id="CMS-2017-0092-1306",
|
|
docket="CMS-2017-0092",
|
|
item_key="K1",
|
|
),
|
|
0.21,
|
|
),
|
|
]
|
|
pool = MagicMock()
|
|
out = retrieve("telehealth", cfg=CFG, pool=pool, k=3)
|
|
assert out == [
|
|
{
|
|
"comment_id": "CMS-2017-0092-1306",
|
|
"docket": "CMS-2017-0092",
|
|
"snippet": "Telehealth comment.",
|
|
"score": 0.21,
|
|
}
|
|
]
|
|
store.similarity_search_with_score.assert_called_once_with("telehealth", k=3)
|
|
|
|
@patch("llm.index.vectorstore")
|
|
def test_falls_back_to_item_key_when_no_comment_id(self, mock_vs):
|
|
store = mock_vs.return_value
|
|
store.similarity_search_with_score.return_value = [
|
|
(_doc("x", item_key="K9", docket="D"), 0.5)
|
|
]
|
|
out = retrieve("q", cfg=CFG, pool=MagicMock())
|
|
assert out[0]["comment_id"] == "K9"
|
|
|
|
|
|
class TestBuildMessages:
|
|
def test_includes_ids_and_abstention_rule(self):
|
|
sources = [
|
|
{
|
|
"comment_id": "CMS-2017-0092-1",
|
|
"docket": "D",
|
|
"snippet": "reduce documentation",
|
|
"score": 0.1,
|
|
}
|
|
]
|
|
msgs = build_messages("why?", sources)
|
|
assert msgs[0]["role"] == "system"
|
|
assert "only" in msgs[0]["content"].lower()
|
|
assert "don't have information" in msgs[0]["content"].lower()
|
|
assert "[CMS-2017-0092-1]" 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 comments" in msgs[1]["content"].lower()
|
|
|
|
|
|
class TestStreamAnswer:
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_yields_tokens_then_sources_then_done(self, mock_retrieve, MockClient):
|
|
mock_retrieve.return_value = [
|
|
{"comment_id": "C1", "docket": "D", "snippet": "s", "score": 0.1}
|
|
]
|
|
lines = [
|
|
'{"message":{"content":"Doc"},"done":false}',
|
|
"", # keep-alive blank line — must be skipped, not parsed
|
|
'{"message":{"content":"tors"},"done":false}',
|
|
'{"message":{"content":""},"done":true}',
|
|
]
|
|
stream_cm = MockClient.return_value.__enter__.return_value.stream.return_value
|
|
resp = stream_cm.__enter__.return_value
|
|
resp.iter_lines.return_value = iter(lines)
|
|
pool = MagicMock()
|
|
|
|
events = list(stream_answer("q", cfg=CFG, pool=pool))
|
|
|
|
pool.check.assert_called_once_with("chat")
|
|
assert events[0] == {"type": "token", "text": "Doc"}
|
|
assert events[1] == {"type": "token", "text": "tors"}
|
|
assert events[-2] == {
|
|
"type": "sources",
|
|
"sources": [
|
|
{"comment_id": "C1", "docket": "D", "snippet": "s", "score": 0.1}
|
|
],
|
|
}
|
|
assert events[-1] == {"type": "done"}
|
|
|
|
@patch("llm.rag.httpx.Client")
|
|
@patch("llm.rag.retrieve")
|
|
def test_http_error_propagates(self, mock_retrieve, MockClient):
|
|
mock_retrieve.return_value = []
|
|
stream_cm = MockClient.return_value.__enter__.return_value.stream.return_value
|
|
resp = stream_cm.__enter__.return_value
|
|
resp.raise_for_status.side_effect = RuntimeError("ollama down")
|
|
with __import__("pytest").raises(RuntimeError, match="ollama down"):
|
|
list(stream_answer("q", cfg=CFG, pool=MagicMock()))
|