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
+101
View File
@@ -0,0 +1,101 @@
"""pgvector cosine retrieval → parent-document mapping (PLAN §3/§6, A7).
Retrieval returns the *chunks* closest to the question embedding (top-K by
cosine distance). The product requirement is that the LLM receives the
**entire relevant document**, not just the chunk (LOCKED A7) — so
:meth:`select_documents` maps chunk hits back to their parent documents
(``chunks.document_id → documents``), dedupes, ranks by best chunk score,
and caps the combined context at ``BOR_MAX_CONTEXT_CHARS``.
"""
from __future__ import annotations
import uuid
from collections.abc import Sequence
from dataclasses import dataclass
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.config import get_settings
from app.models import Chunk, Document
#: Marker appended when the context budget is exceeded (PLAN §6).
TRUNCATION_MARKER = "[…truncated…]"
@dataclass
class RetrievedChunk:
"""One chunk hit: its cosine score plus the parent document row."""
chunk_id: uuid.UUID
position: int
content: str
score: float # 1 − cosine_distance (higher is more similar)
document: Document
def retrieve(
db: Session, question_embedding: list[float], top_k: int | None = None
) -> list[RetrievedChunk]:
"""Top-*top_k* chunks by pgvector cosine distance (``<=>``).
``score = 1 − distance``. Results are ordered by ascending distance, so
index 0 is the best hit. Chunks whose embedding is still NULL (two-phase
import in progress) are skipped.
"""
k = top_k if top_k is not None else get_settings().top_k_chunks
distance = Chunk.embedding.cosine_distance(question_embedding)
rows = db.execute(
select(Chunk, distance.label("distance"), Document)
.join(Document, Chunk.document_id == Document.id)
.where(Chunk.embedding.is_not(None))
.order_by(distance)
.limit(k)
).all()
return [
RetrievedChunk(
chunk_id=chunk.id,
position=chunk.position,
content=chunk.content,
score=round(1.0 - float(dist), 6),
document=doc,
)
for chunk, dist, doc in rows
]
def select_documents(
chunks: Sequence[RetrievedChunk],
n: int | None = None,
max_chars: int | None = None,
) -> list[Document]:
"""Map chunk hits to distinct parent documents, ranked by best chunk score.
At most *n* documents are returned (default ``BOR_TOP_N_DOCS``). The
returned rows carry the full document content; if the combined content
would exceed *max_chars* (default ``BOR_MAX_CONTEXT_CHARS``), the
lowest-ranked overflowing document is truncated in place with the
``[…truncated…]`` marker so the assembled context never exceeds the
budget (PLAN §6).
"""
top_n = n if n is not None else get_settings().top_n_docs
budget = max_chars if max_chars is not None else get_settings().max_context_chars
docs: list[Document] = []
seen: set[uuid.UUID] = set()
for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
if rc.document.id in seen:
continue
seen.add(rc.document.id)
docs.append(rc.document)
docs = docs[:top_n]
remaining = budget
for doc in docs:
if len(doc.content) <= remaining:
remaining -= len(doc.content)
else:
keep = max(0, remaining - len(TRUNCATION_MARKER))
doc.content = doc.content[:keep] + TRUNCATION_MARKER
remaining = 0
return docs