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:
2026-08-21 17:17:02 -04:00
parent 99c48cbe06
commit 396e4d47fb
14 changed files with 1266 additions and 59 deletions
+97 -1
View File
@@ -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"}]))
+88
View File
@@ -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")
+96
View File
@@ -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) == []
+48
View File
@@ -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"