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:
+152
-20
@@ -1,29 +1,161 @@
|
||||
"""POST /api/chat — placeholder (phase 01).
|
||||
"""POST /api/chat — a RAG chat turn streamed over SSE (PLAN §3/§4).
|
||||
|
||||
Phase 03 replaces this with the real RAG pipeline and SSE streaming
|
||||
(PLAN §4 contract: ``delta`` events + final ``done``). The placeholder
|
||||
keeps the same JSON shape the frontend already consumes, so the UI round-trip
|
||||
is exercised end-to-end from day one.
|
||||
Flow (LOCKED A7/A15): embed the question → pgvector cosine top-K chunks →
|
||||
distinct parent documents (full text, capped) → locked persona prompt
|
||||
(PLAN §6) → ``turbo`` streamed as ``delta`` events → final ``done`` event
|
||||
(``deflected``, ``sources``, ``suggestions``) + ``query_log`` row + the
|
||||
per-turn log line (PLAN §9). Mid-stream failures become a structured
|
||||
``error`` event; a pre-stream DB outage is a plain 503 JSON.
|
||||
|
||||
The honesty gate (LOW relevance → deflection) lands in phase 04; every
|
||||
turn in this phase is grounded (``deflected=false``).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from app.schemas import ChatRequest
|
||||
from fastapi import APIRouter, Depends
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
router = APIRouter()
|
||||
from app.config import get_settings
|
||||
from app.db import db_available, get_db
|
||||
from app.models import QueryLog
|
||||
from app.rag.llm import EmbeddingError, LLMClient, LLMError
|
||||
from app.rag.prompts import build_high_prompt
|
||||
from app.rag.retriever import retrieve, select_documents
|
||||
from app.schemas import ChatDoneEvent, ChatRequest, SourceRef
|
||||
|
||||
logger = logging.getLogger("app.chat")
|
||||
router = APIRouter(tags=["chat"])
|
||||
|
||||
#: Streaming hints: no proxy buffering, no client caching (PLAN A15).
|
||||
SSE_HEADERS = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}
|
||||
|
||||
_llm: LLMClient | None = None
|
||||
|
||||
|
||||
def get_llm() -> LLMClient:
|
||||
"""Shared LLM client (FastAPI dependency so tests can override it)."""
|
||||
global _llm
|
||||
if _llm is None:
|
||||
_llm = LLMClient(get_settings())
|
||||
return _llm
|
||||
|
||||
|
||||
def sse_event(payload: dict[str, Any]) -> str:
|
||||
"""Serialize one SSE frame: ``data: <json>\\n\\n`` (PLAN §4)."""
|
||||
return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
||||
|
||||
|
||||
@router.post("/chat")
|
||||
async def chat(request: ChatRequest) -> dict[str, object]:
|
||||
"""Placeholder answer — no LLM, no DB."""
|
||||
return {
|
||||
"ok": True,
|
||||
"answer": (
|
||||
"Hey! My neurons are still wiring up — the real Brain "
|
||||
"(RAG over your docs, powered by aipi) lands in the next "
|
||||
"phases. Try me again soon! 🧠"
|
||||
),
|
||||
"deflected": False,
|
||||
"sources": [],
|
||||
}
|
||||
async def chat(
|
||||
request: ChatRequest,
|
||||
db: Session = Depends(get_db), # noqa: B008
|
||||
llm: LLMClient = Depends(get_llm), # noqa: B008
|
||||
):
|
||||
"""One chat turn: SSE stream of ``delta`` events + a final ``done``."""
|
||||
if not db_available():
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
content={
|
||||
"detail": (
|
||||
"The knowledge base is offline — start Postgres with "
|
||||
"`podman compose up -d db`, then ask again."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
started = time.monotonic()
|
||||
|
||||
async def stream() -> AsyncIterator[str]:
|
||||
# 1. Embed the question.
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
question_vec = await llm.embed_one(request.message)
|
||||
except EmbeddingError as e:
|
||||
logger.error("chat: embedding failed question=%r — %s", request.message, e)
|
||||
yield sse_event(
|
||||
{
|
||||
"type": "error",
|
||||
"detail": "I couldn't reach the embedding model — please try again.",
|
||||
}
|
||||
)
|
||||
return
|
||||
embed_ms = int((time.monotonic() - t0) * 1000)
|
||||
|
||||
# 2. Retrieve top-K chunks → top-N full parent documents.
|
||||
try:
|
||||
chunks = retrieve(db, question_vec)
|
||||
docs = select_documents(chunks)
|
||||
except Exception: # noqa: BLE001 — DB failure mid-turn
|
||||
logger.exception("chat: retrieval failed question=%r", request.message)
|
||||
yield sse_event(
|
||||
{
|
||||
"type": "error",
|
||||
"detail": "The knowledge base went offline mid-question — is Postgres up?",
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
top_score = chunks[0].score if chunks else 0.0
|
||||
source_paths = [f"{d.source}/{d.path}" for d in docs]
|
||||
messages = [
|
||||
{"role": "system", "content": build_high_prompt(docs)},
|
||||
{"role": "user", "content": request.message},
|
||||
]
|
||||
|
||||
# 3. Stream the grounded answer.
|
||||
try:
|
||||
async for piece in llm.chat_stream(messages):
|
||||
yield sse_event({"type": "delta", "text": piece})
|
||||
except LLMError as e:
|
||||
logger.error("chat: LLM stream failed question=%r — %s", request.message, e)
|
||||
yield sse_event(
|
||||
{"type": "error", "detail": "The chat model dropped the connection — try again?"}
|
||||
)
|
||||
return
|
||||
|
||||
# 4. Durable record + required per-turn log line (PLAN §9).
|
||||
total_ms = int((time.monotonic() - started) * 1000)
|
||||
try:
|
||||
db.add(
|
||||
QueryLog(
|
||||
question=request.message,
|
||||
top_score=top_score,
|
||||
chunk_hits=len(chunks),
|
||||
deflected=False,
|
||||
sources=", ".join(source_paths),
|
||||
latency_ms=total_ms,
|
||||
)
|
||||
)
|
||||
db.commit()
|
||||
except Exception: # noqa: BLE001 — the answer already went out
|
||||
logger.exception("chat: failed to write query_log question=%r", request.message)
|
||||
|
||||
logger.info(
|
||||
"question=%r embed_ms=%d top_score=%.3f threshold=%.2f deflected=%s "
|
||||
"sources=%r total_ms=%d",
|
||||
request.message,
|
||||
embed_ms,
|
||||
top_score,
|
||||
get_settings().relevance_threshold,
|
||||
False,
|
||||
source_paths,
|
||||
total_ms,
|
||||
)
|
||||
yield sse_event(
|
||||
ChatDoneEvent(
|
||||
deflected=False,
|
||||
sources=[
|
||||
SourceRef(source=d.source, path=d.path, title=d.title) for d in docs
|
||||
],
|
||||
suggestions=[],
|
||||
).model_dump()
|
||||
)
|
||||
|
||||
return StreamingResponse(stream(), media_type="text/event-stream", headers=SSE_HEADERS)
|
||||
|
||||
+39
-3
@@ -1,7 +1,7 @@
|
||||
"""Async OpenAI-compatible client for the self-hosted aipi endpoint (PLAN A5).
|
||||
|
||||
Phase 02 adds the embeddings surface (the importer — and, from phase 03,
|
||||
retrieval — need it). Chat streaming lands in phase 03 on this same client.
|
||||
Provides the embeddings surface (importer, retrieval) and chat streaming
|
||||
(PLAN A15) for the RAG pipeline.
|
||||
|
||||
Fail-loud rule (PLAN A6): the ``chunks.embedding`` column is fixed at 768
|
||||
dimensions when the table is created, so a model that returns a different
|
||||
@@ -11,8 +11,11 @@ vectors that pgvector rejects.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import cast
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
|
||||
from app.config import Settings, get_settings
|
||||
|
||||
@@ -27,6 +30,10 @@ class EmbeddingDimensionError(EmbeddingError):
|
||||
"""Embedding dimension != BOR_EMBEDDING_DIM — import must fail loudly."""
|
||||
|
||||
|
||||
class LLMError(RuntimeError):
|
||||
"""The chat-completions endpoint failed (network, HTTP, or mid-stream)."""
|
||||
|
||||
|
||||
# aipi's local embedding model rejects requests over ~1024 input tokens
|
||||
# ("input is too large to process"). Batch by estimated tokens, with a
|
||||
# safety margin under that cap — code-dense text can tokenize at ~3
|
||||
@@ -157,6 +164,35 @@ class LLMClient:
|
||||
return out
|
||||
|
||||
async def embed_one(self, text: str) -> list[float]:
|
||||
"""Convenience: embed a single text (retrieval path, phase 03)."""
|
||||
"""Convenience: embed a single text (retrieval path)."""
|
||||
(vec,) = await self.embed([text])
|
||||
return vec
|
||||
|
||||
async def chat_stream(self, messages: list[dict[str, str]]) -> AsyncIterator[str]:
|
||||
"""Stream assistant text deltas from the chat model (PLAN A5/A15).
|
||||
|
||||
``stream=True`` against the OpenAI-compatible endpoint; yields only
|
||||
non-empty ``delta.content`` pieces. Any failure (network, HTTP,
|
||||
malformed stream) surfaces as :class:`LLMError` so the API layer can
|
||||
turn it into an SSE ``error`` event instead of a hung request.
|
||||
"""
|
||||
try:
|
||||
# ``{role, content}`` dicts are exactly what the message params
|
||||
# accept; the cast keeps pyright honest about the SDK's union.
|
||||
stream = await self._client.chat.completions.create(
|
||||
model=self.settings.llm_chat_model,
|
||||
messages=cast("list[ChatCompletionMessageParam]", messages),
|
||||
temperature=0.4,
|
||||
max_tokens=700,
|
||||
stream=True,
|
||||
)
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
piece = chunk.choices[0].delta.content
|
||||
if piece:
|
||||
yield piece
|
||||
except LLMError:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 — wrap transport-level failures
|
||||
raise LLMError(f"chat stream from {self.settings.llm_base_url} failed: {e}") from e
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Locked system-prompt builder (PLAN §6).
|
||||
|
||||
The persona + HONESTY GATE text is **locked verbatim** — change it through
|
||||
the plan, not here. Two modes:
|
||||
|
||||
* ``HIGH`` — grounded turn: full top-document texts under ``<documents>``.
|
||||
* ``LOW`` — deflection turn: weak-hit *titles only* plus the
|
||||
``DEFLECT_MODE`` marker (the E2E mock LLM keys on that marker).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from app.models import Document
|
||||
|
||||
#: PLAN §6 verbatim (line wrapping included); ``{relevance}`` is filled by
|
||||
#: :func:`_base`.
|
||||
PERSONA: str = (
|
||||
'You are "Brain of Reese" — the digital brain of Reese, a self-hoster and\n'
|
||||
"homelab tinkerer. Personality: chippy, upbeat, warm, and genuinely\n"
|
||||
'optimistic about the user\'s ability to do things ("you\'ve got this").\n'
|
||||
"\n"
|
||||
"Rules:\n"
|
||||
"1. Answer ONLY from the provided document context. Cite which document(s)\n"
|
||||
" you used, by path.\n"
|
||||
"2. Be concrete: names, versions, ports, hosts, schedules — the specifics in\n"
|
||||
" the docs are the value.\n"
|
||||
'3. HONESTY GATE: if <relevance> is "LOW", you must NOT pretend to know.\n'
|
||||
' Start your answer with a variant of: "I haven\'t done anything like that."\n'
|
||||
" Then offer 2-3 alternative questions about things you DO have notes on.\n"
|
||||
"4. Never invent facts, hosts, or steps that are not in the context.\n"
|
||||
"5. Keep answers tight: short paragraphs, bullets where helpful.\n"
|
||||
"\n"
|
||||
"<relevance>{relevance}</relevance>"
|
||||
)
|
||||
|
||||
|
||||
def _base(relevance: str) -> str:
|
||||
if relevance not in ("HIGH", "LOW"):
|
||||
raise ValueError(f"relevance must be HIGH or LOW, got {relevance!r}")
|
||||
return PERSONA.replace("{relevance}", relevance)
|
||||
|
||||
|
||||
def build_high_prompt(documents: Sequence[Document]) -> str:
|
||||
"""Grounded turn: locked persona + full texts of the top documents."""
|
||||
blocks = [
|
||||
f'<document source="{doc.source}" path="{doc.path}" title="{doc.title}">\n'
|
||||
f"{doc.content}\n"
|
||||
"</document>"
|
||||
for doc in documents
|
||||
]
|
||||
body = "\n\n".join(blocks) if blocks else (
|
||||
"(no documents matched — do not invent specifics)"
|
||||
)
|
||||
return _base("HIGH") + "\n<documents>\n" + body + "\n</documents>"
|
||||
|
||||
|
||||
def build_deflect_prompt(titles: Sequence[str]) -> str:
|
||||
"""Deflection turn: weak-hit titles only (no document content)."""
|
||||
weak = "\n".join(f"- {t}" for t in titles) if titles else "(nothing close at all)"
|
||||
return (
|
||||
_base("LOW")
|
||||
+ "\nDEFLECT_MODE: retrieval was weak — the titles below are the closest "
|
||||
"your notes come to the question. They are titles only; do not pretend "
|
||||
"they answer it. Use them to propose 2-3 alternative questions.\n"
|
||||
+ weak
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user