feat(rag): stream grounded RAG answers over SSE with source citations
Phase 03 (Story: Chat RAG Answer — happy path):
- app/rag/retriever.py: top-k cosine search + parent-doc selection with
per-doc dedupe and BOR_MAX_CONTEXT_CHARS cap ([…truncated…] marker)
- app/rag/prompts.py: locked persona + HIGH/DEFLECT prompt builders
- app/rag/llm.py: LLMError + chat_stream (turbo, temp 0.4, max 700, stream)
- app/api/chat.py: POST /api/chat SSE — delta* then done{deflected,
sources, suggestions}; query_log row + PLAN §9 per-turn log line;
structured error event on mid-stream failure, JSON 503 when DB down
- frontend: SSE reader, live bubble streaming, source chips -> /sources.html,
red role=alert banner, Send button state that always recovers
- fix(scaffold): [hidden] { display: none !important } — .kb-banner's
display:flex was overriding the hidden attribute (banner always visible)
- tests: unit (retriever/prompts/sse/llm) + integration (real Postgres RAG
turn, query_log, error + 503 paths, mid-turn failures) + Playwright story
suite (grounded answer, log row, raw SSE shape); smoke placeholder test
replaced with the real never-stale-button contract
This commit is contained in:
@@ -10,12 +10,13 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from app.config import Settings
|
||||
from app.rag.llm import EmbeddingDimensionError, EmbeddingError, LLMClient
|
||||
from app.rag.llm import EmbeddingDimensionError, EmbeddingError, LLMClient, LLMError
|
||||
|
||||
|
||||
def _settings(**kwargs: Any) -> Settings:
|
||||
@@ -230,3 +231,98 @@ def test_single_oversized_text_fails_actionably() -> None:
|
||||
with pytest.raises(EmbeddingError, match="token cap"):
|
||||
asyncio.run(llm.embed(["x" * 3000]))
|
||||
assert llm.embed_batches == 0
|
||||
|
||||
|
||||
# ---------- chat streaming (phase 03) ----------
|
||||
|
||||
|
||||
def _chunk(content: str | None = "text", empty: bool = False):
|
||||
"""One fake ChatCompletionChunk (``choices[].delta.content`` shape)."""
|
||||
if empty:
|
||||
return SimpleNamespace(choices=[])
|
||||
return SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=content))])
|
||||
|
||||
|
||||
class _FakeChatStream:
|
||||
def __init__(self, chunks: list) -> None:
|
||||
self._chunks = list(chunks)
|
||||
|
||||
def __aiter__(self):
|
||||
self._i = 0
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._i >= len(self._chunks):
|
||||
raise StopAsyncIteration
|
||||
chunk = self._chunks[self._i]
|
||||
self._i += 1
|
||||
return chunk
|
||||
|
||||
|
||||
class _FakeCompletions:
|
||||
def __init__(self, chunks: list | None = None, fail: Exception | None = None) -> None:
|
||||
self.chunks = chunks or []
|
||||
self.fail = fail
|
||||
self.kwargs: dict | None = None
|
||||
|
||||
async def create(self, **kwargs) -> _FakeChatStream:
|
||||
self.kwargs = kwargs
|
||||
if self.fail is not None:
|
||||
raise self.fail
|
||||
return _FakeChatStream(self.chunks)
|
||||
|
||||
|
||||
def _make_stream_client(
|
||||
chunks: list | None = None, fail: Exception | None = None
|
||||
) -> tuple[LLMClient, _FakeCompletions]:
|
||||
completions = _FakeCompletions(chunks, fail)
|
||||
fake_openai = SimpleNamespace(chat=SimpleNamespace(completions=completions))
|
||||
llm = LLMClient(_settings())
|
||||
llm._client = fake_openai # pyright: ignore[reportAttributeAccessIssue]
|
||||
return llm, completions
|
||||
|
||||
|
||||
async def _collect(llm: LLMClient, messages: list[dict[str, str]]) -> list[str]:
|
||||
return [p async for p in llm.chat_stream(messages)]
|
||||
|
||||
|
||||
def test_chat_stream_yields_deltas_in_order() -> None:
|
||||
llm, completions = _make_stream_client(
|
||||
[_chunk("Hey "), _chunk("you've "), _chunk("got this! 🧠")]
|
||||
)
|
||||
pieces = asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
||||
assert pieces == ["Hey ", "you've ", "got this! 🧠"]
|
||||
|
||||
|
||||
def test_chat_stream_uses_locked_generation_params() -> None:
|
||||
llm, completions = _make_stream_client([_chunk("x")])
|
||||
messages = [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}]
|
||||
asyncio.run(_collect(llm, messages))
|
||||
assert completions.kwargs is not None
|
||||
assert completions.kwargs["model"] == "turbo"
|
||||
assert completions.kwargs["stream"] is True
|
||||
assert completions.kwargs["temperature"] == 0.4
|
||||
assert completions.kwargs["max_tokens"] == 700
|
||||
assert completions.kwargs["messages"] == messages
|
||||
|
||||
|
||||
def test_chat_stream_skips_empty_deltas_and_choiceless_chunks() -> None:
|
||||
llm, _ = _make_stream_client([_chunk("a"), _chunk(empty=True), _chunk(None), _chunk("b")])
|
||||
assert asyncio.run(_collect(llm, [{"role": "user", "content": "q"}])) == ["a", "b"]
|
||||
|
||||
|
||||
def test_chat_stream_wraps_failures_as_llm_error() -> None:
|
||||
llm, _ = _make_stream_client(fail=RuntimeError("connection reset by peer"))
|
||||
|
||||
async def drain() -> None:
|
||||
async for _ in llm.chat_stream([{"role": "user", "content": "q"}]):
|
||||
pass
|
||||
|
||||
with pytest.raises(LLMError, match="connection reset by peer"):
|
||||
asyncio.run(drain())
|
||||
|
||||
|
||||
def test_chat_stream_llm_error_passes_through_unwrapped() -> None:
|
||||
llm, _ = _make_stream_client(fail=LLMError("already wrapped"))
|
||||
with pytest.raises(LLMError, match="already wrapped"):
|
||||
asyncio.run(_collect(llm, [{"role": "user", "content": "q"}]))
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Unit: locked persona prompt builder (PLAN §6 verbatim + both modes)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from app.models import Document
|
||||
from app.rag.prompts import PERSONA, _base, build_deflect_prompt, build_high_prompt
|
||||
|
||||
|
||||
def _doc(path: str, content: str, title: str) -> Document:
|
||||
return Document(
|
||||
id=uuid.uuid4(),
|
||||
source="Homelab",
|
||||
path=path,
|
||||
full_path=f"/tmp/{path}",
|
||||
title=title,
|
||||
content=content,
|
||||
content_hash="0" * 64,
|
||||
)
|
||||
|
||||
|
||||
def test_persona_rules_present_verbatim() -> None:
|
||||
for fragment in (
|
||||
'You are "Brain of Reese" — the digital brain of Reese, a self-hoster and',
|
||||
'optimistic about the user\'s ability to do things ("you\'ve got this")',
|
||||
"Answer ONLY from the provided document context. Cite which document(s)",
|
||||
"you used, by path.",
|
||||
"Be concrete: names, versions, ports, hosts, schedules",
|
||||
'HONESTY GATE: if <relevance> is "LOW", you must NOT pretend to know',
|
||||
'Start your answer with a variant of: "I haven\'t done anything like that."',
|
||||
"Then offer 2-3 alternative questions about things you DO have notes on.",
|
||||
"Never invent facts, hosts, or steps that are not in the context.",
|
||||
"Keep answers tight: short paragraphs, bullets where helpful.",
|
||||
):
|
||||
assert fragment in PERSONA
|
||||
|
||||
|
||||
def test_high_prompt_carries_relevance_marker_and_full_documents() -> None:
|
||||
doc = _doc("kubernetes.md", "Talos Linux on three nodes.", "Kubernetes Homelab Cluster")
|
||||
prompt = build_high_prompt([doc])
|
||||
assert "<relevance>HIGH</relevance>" in prompt
|
||||
assert "DEFLECT_MODE" not in prompt
|
||||
assert "<documents>" in prompt and "</documents>" in prompt
|
||||
assert 'path="kubernetes.md"' in prompt
|
||||
assert "Talos Linux on three nodes." in prompt
|
||||
assert "HONESTY GATE" in prompt # persona intact
|
||||
|
||||
|
||||
def test_high_prompt_lists_multiple_documents_in_order() -> None:
|
||||
a = _doc("a.md", "CONTENT_A", "Title A")
|
||||
b = _doc("b.md", "CONTENT_B", "Title B")
|
||||
prompt = build_high_prompt([a, b])
|
||||
assert prompt.index("CONTENT_A") < prompt.index("CONTENT_B")
|
||||
assert 'title="Title B"' in prompt
|
||||
|
||||
|
||||
def test_high_prompt_without_documents_stays_honest() -> None:
|
||||
prompt = build_high_prompt([])
|
||||
assert "<documents>" in prompt
|
||||
assert "do not invent specifics" in prompt
|
||||
|
||||
|
||||
def test_low_prompt_has_deflect_mode_and_titles_only() -> None:
|
||||
titles = ["Kubernetes Homelab Cluster", "Backup Strategy"]
|
||||
prompt = build_deflect_prompt(titles)
|
||||
assert "<relevance>LOW</relevance>" in prompt
|
||||
assert "DEFLECT_MODE" in prompt # marker the E2E mock keys on
|
||||
assert "- Kubernetes Homelab Cluster" in prompt
|
||||
assert "- Backup Strategy" in prompt
|
||||
|
||||
|
||||
def test_low_prompt_never_contains_document_content() -> None:
|
||||
secret = "SECRET_DOCUMENT_CONTENT_12345"
|
||||
prompt = build_deflect_prompt(["Some Title"])
|
||||
assert secret not in prompt
|
||||
assert "<documents>" not in prompt
|
||||
assert "HONESTY GATE" in prompt # the LOW rule is what the model must follow
|
||||
|
||||
|
||||
def test_low_prompt_with_no_titles() -> None:
|
||||
assert "nothing close at all" in build_deflect_prompt([])
|
||||
|
||||
|
||||
def test_relevance_placeholder_rejected_for_garbage() -> None:
|
||||
with pytest.raises(ValueError, match="HIGH or LOW"):
|
||||
_base("MEDIUM")
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Unit: retriever — ordering, dedup, and the context cap (fake rows).
|
||||
|
||||
The SQL side of :func:`app.rag.retriever.retrieve` is exercised by the
|
||||
chat integration tests against real Postgres; the pure mapping logic in
|
||||
:func:`select_documents` is tested here with in-memory rows.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
from app.models import Document
|
||||
from app.rag.retriever import TRUNCATION_MARKER, RetrievedChunk, select_documents
|
||||
|
||||
|
||||
def _doc(path: str, content: str, source: str = "Homelab", title: str | None = None) -> Document:
|
||||
return Document(
|
||||
id=uuid.uuid4(),
|
||||
source=source,
|
||||
path=path,
|
||||
full_path=f"/tmp/{path}",
|
||||
title=title or path,
|
||||
content=content,
|
||||
content_hash="0" * 64,
|
||||
)
|
||||
|
||||
|
||||
def _chunk(doc: Document, score: float, position: int = 0) -> RetrievedChunk:
|
||||
return RetrievedChunk(
|
||||
chunk_id=uuid.uuid4(),
|
||||
position=position,
|
||||
content=doc.content[:40],
|
||||
score=score,
|
||||
document=doc,
|
||||
)
|
||||
|
||||
|
||||
def test_ranks_by_best_chunk_score_not_first_hit() -> None:
|
||||
"""A doc whose *later* chunk scores highest must still rank first."""
|
||||
a = _doc("a.md", "A" * 50)
|
||||
b = _doc("b.md", "B" * 50)
|
||||
c = _doc("c.md", "C" * 50)
|
||||
chunks = [
|
||||
_chunk(a, 0.4, position=0), # a's weak chunk comes first
|
||||
_chunk(b, 0.8),
|
||||
_chunk(a, 0.9, position=2), # a's best chunk comes last
|
||||
_chunk(c, 0.5),
|
||||
]
|
||||
docs = select_documents(chunks, n=3, max_chars=10_000)
|
||||
assert [d.path for d in docs] == ["a.md", "b.md", "c.md"]
|
||||
|
||||
|
||||
def test_dedups_to_one_document_per_hit_set() -> None:
|
||||
a = _doc("a.md", "A" * 50)
|
||||
chunks = [_chunk(a, 0.2), _chunk(a, 0.7), _chunk(a, 0.5)]
|
||||
docs = select_documents(chunks, n=2, max_chars=10_000)
|
||||
assert len(docs) == 1
|
||||
assert docs[0] is a
|
||||
|
||||
|
||||
def test_caps_at_n_documents() -> None:
|
||||
docs_in = [_doc(f"d{i}.md", "X" * 20) for i in range(4)]
|
||||
chunks = [_chunk(d, 0.5 - 0.1 * i) for i, d in enumerate(docs_in)]
|
||||
out = select_documents(chunks, n=2, max_chars=10_000)
|
||||
assert [d.path for d in out] == ["d0.md", "d1.md"]
|
||||
|
||||
|
||||
def test_combined_content_capped_with_truncation_marker() -> None:
|
||||
big = _doc("big.md", "B" * 100)
|
||||
small = _doc("small.md", "S" * 100)
|
||||
chunks = [_chunk(big, 0.9), _chunk(small, 0.6)]
|
||||
out = select_documents(chunks, n=2, max_chars=150)
|
||||
# Best doc stays intact; the overflowing one is truncated in place.
|
||||
assert out[0].content == "B" * 100
|
||||
assert out[1].content.endswith(TRUNCATION_MARKER)
|
||||
assert out[1].content.startswith("S")
|
||||
assert len(out[0].content) + len(out[1].content) <= 150
|
||||
|
||||
|
||||
def test_single_doc_over_budget_is_truncated_to_budget() -> None:
|
||||
big = _doc("big.md", "Z" * 200)
|
||||
out = select_documents([_chunk(big, 0.9)], n=2, max_chars=50)
|
||||
assert len(out[0].content) == 50
|
||||
assert out[0].content.endswith(TRUNCATION_MARKER)
|
||||
|
||||
|
||||
def test_under_budget_no_truncation() -> None:
|
||||
a = _doc("a.md", "A" * 80)
|
||||
b = _doc("b.md", "B" * 60)
|
||||
out = select_documents([_chunk(b, 0.5), _chunk(a, 0.9)], n=2, max_chars=200)
|
||||
assert [d.path for d in out] == ["a.md", "b.md"]
|
||||
assert a.content == "A" * 80 and b.content == "B" * 60
|
||||
assert TRUNCATION_MARKER not in a.content + b.content
|
||||
|
||||
|
||||
def test_empty_hits_yield_no_documents() -> None:
|
||||
assert select_documents([], n=2, max_chars=24_000) == []
|
||||
@@ -0,0 +1,48 @@
|
||||
"""Unit: SSE frame serialization for POST /api/chat (PLAN §4 contract)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from app.api.chat import sse_event
|
||||
|
||||
|
||||
def _payload(frame: str) -> dict:
|
||||
assert frame.startswith("data: ")
|
||||
assert frame.endswith("\n\n")
|
||||
return json.loads(frame.removeprefix("data: ").strip())
|
||||
|
||||
|
||||
def test_delta_frame_serializes_exactly() -> None:
|
||||
frame = sse_event({"type": "delta", "text": "hi"})
|
||||
assert frame == 'data: {"type": "delta", "text": "hi"}\n\n'
|
||||
assert _payload(frame) == {"type": "delta", "text": "hi"}
|
||||
|
||||
|
||||
def test_done_frame_carries_full_contract_shape() -> None:
|
||||
payload = {
|
||||
"type": "done",
|
||||
"deflected": False,
|
||||
"sources": [{"source": "Homelab", "path": "kubernetes.md", "title": "K8s"}],
|
||||
"suggestions": [],
|
||||
}
|
||||
assert _payload(sse_event(payload)) == payload
|
||||
|
||||
|
||||
def test_error_frame_serializes() -> None:
|
||||
frame = sse_event({"type": "error", "detail": "boom"})
|
||||
assert _payload(frame) == {"type": "error", "detail": "boom"}
|
||||
|
||||
|
||||
def test_unicode_survives_roundtrip() -> None:
|
||||
frame = sse_event({"type": "delta", "text": "🧠 café — \"quoted\""})
|
||||
# ensure_ascii=False keeps the frame readable (no \uXXXX escapes).
|
||||
assert "🧠 café" in frame
|
||||
assert _payload(frame)["text"] == "🧠 café — \"quoted\""
|
||||
|
||||
|
||||
def test_multi_line_text_stays_one_frame() -> None:
|
||||
"""Newlines inside the JSON payload must be escaped so the frame
|
||||
delimiter ``\\n\\n`` remains unambiguous."""
|
||||
frame = sse_event({"type": "delta", "text": "line1\nline2\n\n"})
|
||||
assert frame.count("\n\n") == 1 # only the frame terminator
|
||||
assert _payload(frame)["text"] == "line1\nline2\n\n"
|
||||
Reference in New Issue
Block a user