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
|
Flow (LOCKED A7/A15): embed the question → pgvector cosine top-K chunks →
|
||||||
(PLAN §4 contract: ``delta`` events + final ``done``). The placeholder
|
distinct parent documents (full text, capped) → locked persona prompt
|
||||||
keeps the same JSON shape the frontend already consumes, so the UI round-trip
|
(PLAN §6) → ``turbo`` streamed as ``delta`` events → final ``done`` event
|
||||||
is exercised end-to-end from day one.
|
(``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 __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")
|
@router.post("/chat")
|
||||||
async def chat(request: ChatRequest) -> dict[str, object]:
|
async def chat(
|
||||||
"""Placeholder answer — no LLM, no DB."""
|
request: ChatRequest,
|
||||||
return {
|
db: Session = Depends(get_db), # noqa: B008
|
||||||
"ok": True,
|
llm: LLMClient = Depends(get_llm), # noqa: B008
|
||||||
"answer": (
|
):
|
||||||
"Hey! My neurons are still wiring up — the real Brain "
|
"""One chat turn: SSE stream of ``delta`` events + a final ``done``."""
|
||||||
"(RAG over your docs, powered by aipi) lands in the next "
|
if not db_available():
|
||||||
"phases. Try me again soon! 🧠"
|
return JSONResponse(
|
||||||
),
|
status_code=503,
|
||||||
"deflected": False,
|
content={
|
||||||
"sources": [],
|
"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).
|
"""Async OpenAI-compatible client for the self-hosted aipi endpoint (PLAN A5).
|
||||||
|
|
||||||
Phase 02 adds the embeddings surface (the importer — and, from phase 03,
|
Provides the embeddings surface (importer, retrieval) and chat streaming
|
||||||
retrieval — need it). Chat streaming lands in phase 03 on this same client.
|
(PLAN A15) for the RAG pipeline.
|
||||||
|
|
||||||
Fail-loud rule (PLAN A6): the ``chunks.embedding`` column is fixed at 768
|
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
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
from typing import cast
|
||||||
|
|
||||||
from openai import AsyncOpenAI
|
from openai import AsyncOpenAI
|
||||||
|
from openai.types.chat import ChatCompletionMessageParam
|
||||||
|
|
||||||
from app.config import Settings, get_settings
|
from app.config import Settings, get_settings
|
||||||
|
|
||||||
@@ -27,6 +30,10 @@ class EmbeddingDimensionError(EmbeddingError):
|
|||||||
"""Embedding dimension != BOR_EMBEDDING_DIM — import must fail loudly."""
|
"""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
|
# aipi's local embedding model rejects requests over ~1024 input tokens
|
||||||
# ("input is too large to process"). Batch by estimated tokens, with a
|
# ("input is too large to process"). Batch by estimated tokens, with a
|
||||||
# safety margin under that cap — code-dense text can tokenize at ~3
|
# safety margin under that cap — code-dense text can tokenize at ~3
|
||||||
@@ -157,6 +164,35 @@ class LLMClient:
|
|||||||
return out
|
return out
|
||||||
|
|
||||||
async def embed_one(self, text: str) -> list[float]:
|
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])
|
(vec,) = await self.embed([text])
|
||||||
return vec
|
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
|
||||||
+114
-12
@@ -1,8 +1,11 @@
|
|||||||
/* Brain of Reese — chat shell.
|
/* Brain of Reese — chat shell.
|
||||||
*
|
*
|
||||||
* Scaffolding-stage behavior: renders suggestions, shows KB health, and
|
* Renders suggestions, shows KB health, and runs chat turns against
|
||||||
* echoes a friendly placeholder answer. The real RAG streaming chat is
|
* POST /api/chat (SSE, PLAN §4): deltas render live into the Brain bubble,
|
||||||
* implemented in the chat-rag phase (see .agent/user_stories/).
|
* the done event appends source chips, errors surface as a red banner.
|
||||||
|
* The full feedback state machine lands with the loading-feedback story;
|
||||||
|
* this keeps the "never stale" contract: the button is busy for the whole
|
||||||
|
* turn and is always re-enabled at the end.
|
||||||
* All DOM ids match frontend/index.html.
|
* All DOM ids match frontend/index.html.
|
||||||
*/
|
*/
|
||||||
|
|
||||||
@@ -148,6 +151,70 @@ function autoGrow() {
|
|||||||
input.style.height = `${Math.min(input.scrollHeight, 192)}px`;
|
input.style.height = `${Math.min(input.scrollHeight, 192)}px`;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/* ---------- chat turn (SSE streaming, PLAN §4) ---------- */
|
||||||
|
|
||||||
|
/* Parse an SSE response body into JSON events. */
|
||||||
|
async function readSSE(response, onEvent) {
|
||||||
|
const reader = response.body.getReader();
|
||||||
|
const decoder = new TextDecoder();
|
||||||
|
let buf = "";
|
||||||
|
for (;;) {
|
||||||
|
const { done, value } = await reader.read();
|
||||||
|
if (done) break;
|
||||||
|
buf += decoder.decode(value, { stream: true });
|
||||||
|
let sep;
|
||||||
|
while ((sep = buf.indexOf("\n\n")) !== -1) {
|
||||||
|
const frame = buf.slice(0, sep).trim();
|
||||||
|
buf = buf.slice(sep + 2);
|
||||||
|
if (!frame.startsWith("data:")) continue;
|
||||||
|
const payload = frame.slice(5).trim();
|
||||||
|
if (!payload || payload === "[DONE]") continue;
|
||||||
|
onEvent(JSON.parse(payload));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/* Source chips (mono, source/path) under a Brain bubble. */
|
||||||
|
function appendSources(wrap, sources) {
|
||||||
|
if (!sources || !sources.length) return;
|
||||||
|
const body = wrap.querySelector(".msg-body");
|
||||||
|
const meta = document.createElement("div");
|
||||||
|
meta.className = "msg-meta";
|
||||||
|
meta.setAttribute("role", "list");
|
||||||
|
meta.setAttribute("aria-label", "Sources");
|
||||||
|
for (const s of sources) {
|
||||||
|
const label = `${s.source}/${s.path}`;
|
||||||
|
const chip = document.createElement("a");
|
||||||
|
chip.className = "source-chip";
|
||||||
|
chip.setAttribute("role", "listitem");
|
||||||
|
chip.href = "/sources.html";
|
||||||
|
chip.textContent = label;
|
||||||
|
chip.title = label;
|
||||||
|
meta.appendChild(chip);
|
||||||
|
}
|
||||||
|
body.appendChild(meta);
|
||||||
|
// Accessible full path whenever the pill visually truncates.
|
||||||
|
for (const chip of meta.children) {
|
||||||
|
if (chip.scrollWidth > chip.clientWidth) chip.setAttribute("aria-label", chip.title);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function showErrorBanner(detail) {
|
||||||
|
banner.hidden = false;
|
||||||
|
banner.classList.add("is-error");
|
||||||
|
banner.setAttribute("role", "alert");
|
||||||
|
bannerText.textContent = `${detail} Try your question again — I'm ready.`;
|
||||||
|
}
|
||||||
|
|
||||||
|
function clearErrorBanner() {
|
||||||
|
if (banner.classList.contains("is-error")) {
|
||||||
|
banner.classList.remove("is-error");
|
||||||
|
banner.setAttribute("role", "status");
|
||||||
|
bannerText.textContent = "";
|
||||||
|
banner.hidden = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async function handleSend(e) {
|
async function handleSend(e) {
|
||||||
e.preventDefault();
|
e.preventDefault();
|
||||||
const text = input.value.trim();
|
const text = input.value.trim();
|
||||||
@@ -156,21 +223,56 @@ async function handleSend(e) {
|
|||||||
addMessage("user", renderMarkdown(text));
|
addMessage("user", renderMarkdown(text));
|
||||||
input.value = "";
|
input.value = "";
|
||||||
autoGrow();
|
autoGrow();
|
||||||
|
clearErrorBanner();
|
||||||
setBusy(true);
|
setBusy(true);
|
||||||
addTyping();
|
addTyping();
|
||||||
|
|
||||||
|
let wrap = null;
|
||||||
|
let acc = "";
|
||||||
|
let res = null;
|
||||||
try {
|
try {
|
||||||
// TODO(chat-rag phase): replace with POST /api/chat (SSE streaming).
|
res = await fetch("/api/chat", {
|
||||||
await new Promise((res) => setTimeout(res, 500));
|
method: "POST",
|
||||||
const reply =
|
headers: { "Content-Type": "application/json" },
|
||||||
"I'm still getting my neurons wired up — the real me ships in the " +
|
body: JSON.stringify({ message: text }),
|
||||||
"next phase! Keep the questions coming, you're on a roll. 🚀";
|
});
|
||||||
|
if (!res.ok || !res.body) {
|
||||||
|
let detail = `Brain's API answered with HTTP ${res.status}.`;
|
||||||
|
try {
|
||||||
|
const body = await res.json();
|
||||||
|
if (body.detail) detail = body.detail;
|
||||||
|
} catch { /* non-JSON error body */ }
|
||||||
|
throw new Error(detail);
|
||||||
|
}
|
||||||
|
await readSSE(res, (ev) => {
|
||||||
|
if (ev.type === "delta") {
|
||||||
|
acc += ev.text || "";
|
||||||
|
if (!wrap) {
|
||||||
|
removeTyping();
|
||||||
|
wrap = addMessage("brain", "");
|
||||||
|
}
|
||||||
|
wrap.querySelector(".bubble").innerHTML = renderMarkdown(acc);
|
||||||
|
wrap.scrollIntoView({ behavior: "smooth", block: "end" });
|
||||||
|
} else if (ev.type === "done") {
|
||||||
|
if (!wrap) {
|
||||||
|
removeTyping();
|
||||||
|
wrap = addMessage("brain", "…");
|
||||||
|
}
|
||||||
|
if (ev.deflected) wrap.classList.add("is-deflected");
|
||||||
|
appendSources(wrap, ev.sources);
|
||||||
|
} else if (ev.type === "error") {
|
||||||
|
throw new Error(ev.detail || "Something went wrong on my side.");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if (!wrap) {
|
||||||
|
removeTyping();
|
||||||
|
addMessage("brain", "Hmm — that came back empty. Ask me again?");
|
||||||
|
}
|
||||||
|
} catch (err) {
|
||||||
removeTyping();
|
removeTyping();
|
||||||
addMessage("brain", renderMarkdown(reply));
|
showErrorBanner(err.message || "Something went wrong on my side.");
|
||||||
} catch {
|
|
||||||
removeTyping();
|
|
||||||
addMessage("brain", "Something went wrong on my side — please try again in a moment.");
|
|
||||||
} finally {
|
} finally {
|
||||||
|
try { res?.body?.cancel(); } catch { /* stream already closed */ }
|
||||||
setBusy(false);
|
setBusy(false);
|
||||||
input.focus();
|
input.focus();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -56,6 +56,11 @@ body {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/* ---------- Accessibility helpers ---------- */
|
/* ---------- Accessibility helpers ---------- */
|
||||||
|
/* The `hidden` attribute must always win — some components set an explicit
|
||||||
|
`display` (e.g. .kb-banner { display: flex }) which would otherwise override
|
||||||
|
the UA stylesheet's `[hidden] { display: none }` and leave the element visible. */
|
||||||
|
[hidden] { display: none !important; }
|
||||||
|
|
||||||
.visually-hidden {
|
.visually-hidden {
|
||||||
position: absolute !important;
|
position: absolute !important;
|
||||||
width: 1px; height: 1px;
|
width: 1px; height: 1px;
|
||||||
|
|||||||
@@ -0,0 +1,170 @@
|
|||||||
|
"""Phase 03 E2E (Playwright): the happy-path RAG chat turn.
|
||||||
|
|
||||||
|
Story: ``.agent/user_stories/chat-rag-answer.md``
|
||||||
|
Run in isolation (DB must be up: ``podman compose up -d db``):
|
||||||
|
|
||||||
|
uv run pytest tests/e2e/test_chat_rag.py -v --no-cov
|
||||||
|
|
||||||
|
Seeding reuses the real importer against ``tests/fixtures/docs/`` with the
|
||||||
|
deterministic mock embeddings (same pattern as the phase 02 story suite);
|
||||||
|
the mock LLM answers on-topic questions by quoting the question and the
|
||||||
|
document context, so the UI assertions are fully deterministic.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
from threading import Thread
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from playwright.sync_api import Page, expect
|
||||||
|
from sqlalchemy import select, text
|
||||||
|
|
||||||
|
from app.config import Settings
|
||||||
|
from app.db import SessionLocal
|
||||||
|
from app.models import QueryLog
|
||||||
|
from app.rag.importer import ImportSummary, import_sources
|
||||||
|
from app.rag.llm import LLMClient
|
||||||
|
|
||||||
|
REPO = Path(__file__).resolve().parents[2]
|
||||||
|
FIXTURES = REPO / "tests" / "fixtures" / "docs"
|
||||||
|
QUESTION = "How is my Kubernetes cluster set up?"
|
||||||
|
MOCK_ANSWER_MARKER = "Deterministic mock answer for E2E"
|
||||||
|
|
||||||
|
|
||||||
|
async def _import_fixtures(mock_port: int) -> ImportSummary:
|
||||||
|
kwargs: dict[str, Any] = {"_env_file": None, "llm_base_url": f"http://127.0.0.1:{mock_port}/v1"}
|
||||||
|
settings = Settings(**kwargs) # pyright: ignore[reportCallIssue]
|
||||||
|
return await import_sources([FIXTURES], LLMClient(settings))
|
||||||
|
|
||||||
|
|
||||||
|
def _run_in_thread(coro: Any) -> Any:
|
||||||
|
"""Run a coroutine on a worker thread.
|
||||||
|
|
||||||
|
Playwright's sync API keeps an asyncio loop running on the test thread,
|
||||||
|
so ``asyncio.run`` cannot be called directly from a test body.
|
||||||
|
"""
|
||||||
|
box: dict[str, Any] = {}
|
||||||
|
|
||||||
|
def runner() -> None:
|
||||||
|
try:
|
||||||
|
box["value"] = asyncio.run(coro)
|
||||||
|
except BaseException as e: # noqa: BLE001 — re-raised on the test thread
|
||||||
|
box["error"] = e
|
||||||
|
|
||||||
|
t = Thread(target=runner)
|
||||||
|
t.start()
|
||||||
|
t.join()
|
||||||
|
if "error" in box:
|
||||||
|
raise box["error"]
|
||||||
|
return box["value"]
|
||||||
|
|
||||||
|
|
||||||
|
def _reset_db(mock_port: int, seed: bool) -> ImportSummary | None:
|
||||||
|
"""Truncate the KB (and query log), then optionally re-import fixtures."""
|
||||||
|
with SessionLocal() as db:
|
||||||
|
db.execute(text("TRUNCATE chunks, documents, query_log"))
|
||||||
|
db.commit()
|
||||||
|
if not seed:
|
||||||
|
return None
|
||||||
|
return _run_in_thread(_import_fixtures(mock_port))
|
||||||
|
|
||||||
|
|
||||||
|
def test_on_topic_question_streams_grounded_answer(
|
||||||
|
page: Page, app_url: str, mock_llm: int, db_ready: None
|
||||||
|
) -> None:
|
||||||
|
summary = _reset_db(mock_llm, seed=True)
|
||||||
|
assert summary is not None and summary.added == 3
|
||||||
|
page.set_default_timeout(30_000)
|
||||||
|
page.goto(app_url)
|
||||||
|
|
||||||
|
# Healthy KB: the offline banner must stay hidden.
|
||||||
|
expect(page.locator("#kb-banner")).to_be_hidden()
|
||||||
|
|
||||||
|
page.fill("#message-input", QUESTION)
|
||||||
|
page.click("#send-btn")
|
||||||
|
|
||||||
|
# User bubble (right, brand) shows the question.
|
||||||
|
expect(page.locator(".msg.user .bubble")).to_contain_text(QUESTION)
|
||||||
|
|
||||||
|
# Brain bubble streams in: the mock quotes the question and ends with a
|
||||||
|
# deterministic marker — waiting on the marker proves content arrived.
|
||||||
|
bubble = page.locator(".msg.brain .bubble")
|
||||||
|
bubble.first.wait_for(state="visible", timeout=30_000)
|
||||||
|
expect(bubble.first).to_contain_text(QUESTION, timeout=30_000)
|
||||||
|
expect(bubble.first).to_contain_text(MOCK_ANSWER_MARKER, timeout=30_000)
|
||||||
|
|
||||||
|
# Grounded: a kubernetes.md source chip renders under the bubble
|
||||||
|
# (top-N docs can add more chips; the question's doc must be among them).
|
||||||
|
chip = page.locator(".msg.brain .source-chip", has_text="kubernetes.md")
|
||||||
|
expect(chip).to_have_count(1)
|
||||||
|
expect(chip.first).to_contain_text("kubernetes.md")
|
||||||
|
expect(chip.first).to_have_attribute("href", "/sources.html")
|
||||||
|
|
||||||
|
# Button recovers: enabled + "Send" (never stale).
|
||||||
|
expect(page.locator("#send-btn")).to_be_enabled()
|
||||||
|
expect(page.locator("#send-label")).to_have_text("Send")
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_logs_query(page: Page, app_url: str, mock_llm: int, db_ready: None) -> None:
|
||||||
|
_reset_db(mock_llm, seed=True)
|
||||||
|
page.set_default_timeout(30_000)
|
||||||
|
page.goto(app_url)
|
||||||
|
page.fill("#message-input", QUESTION)
|
||||||
|
page.click("#send-btn")
|
||||||
|
expect(
|
||||||
|
page.locator(".msg.brain .source-chip", has_text="kubernetes.md")
|
||||||
|
).to_have_count(1, timeout=30_000)
|
||||||
|
|
||||||
|
# App still healthy after the turn.
|
||||||
|
r = httpx.get(f"{app_url}/api/health", timeout=5)
|
||||||
|
assert r.status_code == 200
|
||||||
|
body = r.json()
|
||||||
|
assert body["status"] == "ok" and body["db"] == "up"
|
||||||
|
|
||||||
|
# Durable record: exactly one query_log row for the turn.
|
||||||
|
with SessionLocal() as db:
|
||||||
|
rows = db.scalars(select(QueryLog)).all()
|
||||||
|
assert len(rows) == 1
|
||||||
|
row = rows[0]
|
||||||
|
assert row.question == QUESTION
|
||||||
|
assert row.deflected is False
|
||||||
|
assert row.top_score > 0.0
|
||||||
|
assert row.chunk_hits >= 1
|
||||||
|
assert "docs/homelab/kubernetes.md" in row.sources
|
||||||
|
assert row.latency_ms >= 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_sse_stream_shape(app_url: str, mock_llm: int, db_ready: None) -> None:
|
||||||
|
"""Raw transport contract (PLAN §4): delta events, then one done."""
|
||||||
|
_reset_db(mock_llm, seed=True)
|
||||||
|
|
||||||
|
frames: list[dict[str, Any]] = []
|
||||||
|
with httpx.stream(
|
||||||
|
"POST", f"{app_url}/api/chat", json={"message": QUESTION}, timeout=60.0
|
||||||
|
) as r:
|
||||||
|
assert r.status_code == 200
|
||||||
|
assert r.headers["content-type"].startswith("text/event-stream")
|
||||||
|
buf = ""
|
||||||
|
for part in r.iter_text():
|
||||||
|
buf += part
|
||||||
|
while "\n\n" in buf:
|
||||||
|
frame, buf = buf.split("\n\n", 1)
|
||||||
|
if frame.strip().startswith("data:"):
|
||||||
|
frames.append(json.loads(frame.strip().removeprefix("data:").strip()))
|
||||||
|
assert buf.strip() == "" # stream ends cleanly on a frame boundary
|
||||||
|
|
||||||
|
deltas = [f for f in frames if f.get("type") == "delta"]
|
||||||
|
assert len(deltas) >= 2, "answer must arrive as multiple deltas (streamed)"
|
||||||
|
assert all(d.get("text") for d in deltas)
|
||||||
|
assert "".join(d["text"] for d in deltas) # non-empty answer
|
||||||
|
|
||||||
|
done = [f for f in frames if f.get("type") == "done"]
|
||||||
|
assert len(done) == 1
|
||||||
|
assert frames[-1]["type"] == "done" # done is the final event
|
||||||
|
assert done[0]["deflected"] is False
|
||||||
|
assert done[0]["suggestions"] == []
|
||||||
|
assert done[0]["sources"], "done must carry the cited sources"
|
||||||
|
assert any(s["path"] == "homelab/kubernetes.md" for s in done[0]["sources"])
|
||||||
+13
-13
@@ -1,12 +1,10 @@
|
|||||||
"""Phase 01 smoke E2E: the app boots, serves the local frontend, and the
|
"""Phase 01 smoke E2E: the app boots, serves the local frontend, and a chat
|
||||||
placeholder chat round-trips without a stale button.
|
round-trip never leaves a stale button (answer or error banner, both fine).
|
||||||
|
|
||||||
Run: uv run pytest tests/e2e/test_smoke.py -v
|
Run: uv run pytest tests/e2e/test_smoke.py -v
|
||||||
"""
|
"""
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import re
|
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
from playwright.sync_api import Page, expect
|
from playwright.sync_api import Page, expect
|
||||||
|
|
||||||
@@ -27,19 +25,21 @@ def test_index_page_loads_locally(page: Page, app_url: str) -> None:
|
|||||||
assert 'href="http' not in html.replace('href="http://www.w3.org', "")
|
assert 'href="http' not in html.replace('href="http://www.w3.org', "")
|
||||||
|
|
||||||
|
|
||||||
def test_placeholder_chat_roundtrip(page: Page, app_url: str) -> None:
|
def test_chat_roundtrip_never_stale_button(page: Page, app_url: str) -> None:
|
||||||
page.goto(app_url)
|
page.goto(app_url)
|
||||||
page.locator("#message-input").fill("hello brain")
|
page.locator("#message-input").fill("hello brain")
|
||||||
page.locator("#send-btn").click()
|
page.locator("#send-btn").click()
|
||||||
|
|
||||||
# User bubble appears, then the Brain placeholder answer arrives.
|
# User bubble appears first.
|
||||||
page.locator(".msg.user .bubble").first.wait_for(state="visible", timeout=10_000)
|
page.locator(".msg.user .bubble").first.wait_for(state="visible", timeout=10_000)
|
||||||
brain_bubble = page.locator(".msg.brain .bubble").first
|
# Then either a streamed Brain answer (DB up) or an error banner
|
||||||
brain_bubble.wait_for(state="visible", timeout=10_000)
|
# (DB down) — but the turn must always complete.
|
||||||
# to_have_text retries until the async fetch resolves (no stale read).
|
page.wait_for_selector(
|
||||||
expect(brain_bubble).to_have_text(re.compile("neurons"), timeout=10_000)
|
".msg.brain .bubble, #kb-banner.is-error",
|
||||||
|
state="visible",
|
||||||
|
timeout=20_000,
|
||||||
|
)
|
||||||
|
|
||||||
# Button is never left stuck: back to "Send" and enabled.
|
# Button is never left stuck: back to "Send" and enabled.
|
||||||
btn = page.locator("#send-btn")
|
expect(page.locator("#send-btn")).to_be_enabled(timeout=10_000)
|
||||||
assert btn.is_enabled()
|
expect(page.locator("#send-label")).to_have_text("Send", timeout=10_000)
|
||||||
assert "Send" in btn.inner_text()
|
|
||||||
|
|||||||
@@ -34,16 +34,6 @@ def test_styles_and_js_served(client) -> None:
|
|||||||
assert client.get("/assets/app.js").status_code == 200
|
assert client.get("/assets/app.js").status_code == 200
|
||||||
|
|
||||||
|
|
||||||
def test_chat_placeholder_roundtrip(client) -> None:
|
|
||||||
r = client.post("/api/chat", json={"message": "hello brain"})
|
|
||||||
assert r.status_code == 200
|
|
||||||
data = r.json()
|
|
||||||
assert data["ok"] is True
|
|
||||||
assert "neurons" in data["answer"]
|
|
||||||
assert data["deflected"] is False
|
|
||||||
assert data["sources"] == []
|
|
||||||
|
|
||||||
|
|
||||||
def test_chat_requires_message(client) -> None:
|
def test_chat_requires_message(client) -> None:
|
||||||
r = client.post("/api/chat", json={"message": ""})
|
r = client.post("/api/chat", json={"message": ""})
|
||||||
assert r.status_code == 422
|
assert r.status_code == 422
|
||||||
|
|||||||
@@ -0,0 +1,276 @@
|
|||||||
|
"""Integration: POST /api/chat — the RAG turn end-to-end.
|
||||||
|
|
||||||
|
Real Postgres (compose) seeded from ``tests/fixtures/docs/`` through the
|
||||||
|
real importer; the LLM client is a deterministic in-process fake
|
||||||
|
(token-overlap embeddings, canned streamed answer), so no network is
|
||||||
|
needed and the cosine ordering is meaningful: the Kubernetes question
|
||||||
|
retrieves the Kubernetes document.
|
||||||
|
|
||||||
|
Requires: podman compose up -d db
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import re
|
||||||
|
from collections.abc import Iterator
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
from sqlalchemy import func, select, text
|
||||||
|
|
||||||
|
from app.api import chat as chat_api
|
||||||
|
from app.config import Settings, get_settings
|
||||||
|
from app.main import app as fastapi_app
|
||||||
|
from app.models import Chunk, QueryLog
|
||||||
|
from app.rag.importer import import_sources
|
||||||
|
from app.rag.llm import EmbeddingError, LLMError
|
||||||
|
|
||||||
|
FIXTURES = Path(__file__).resolve().parents[1] / "fixtures" / "docs"
|
||||||
|
QUESTION = "How is my Kubernetes cluster set up?"
|
||||||
|
DIM = 768
|
||||||
|
_TOKEN_RE = re.compile(r"[a-z0-9]+")
|
||||||
|
|
||||||
|
|
||||||
|
def _token_vec(text: str) -> list[float]:
|
||||||
|
"""Bag-of-words unit vector — same algorithm as the E2E mock, so the
|
||||||
|
cosine behaviour here matches what the story E2E sees."""
|
||||||
|
vec = [0.0] * DIM
|
||||||
|
for tok in _TOKEN_RE.findall(text.lower()):
|
||||||
|
vec[int(hashlib.md5(tok.encode()).hexdigest(), 16) % DIM] += 1.0
|
||||||
|
norm = math.sqrt(sum(v * v for v in vec)) or 1.0
|
||||||
|
return [v / norm for v in vec]
|
||||||
|
|
||||||
|
|
||||||
|
class FakeRagLLM:
|
||||||
|
"""Duck-typed :class:`app.rag.llm.LLMClient` stand-in for the chat path."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
answer: str = "Hey — you've got this! Talos, Cilium, three nodes. 🧠",
|
||||||
|
embed_error: Exception | None = None,
|
||||||
|
stream_error: Exception | None = None,
|
||||||
|
fail_mid_stream: bool = False,
|
||||||
|
) -> None:
|
||||||
|
self.settings = Settings(_env_file=None) # pyright: ignore[reportCallIssue]
|
||||||
|
self.embed_batches = 0
|
||||||
|
self.answer = answer
|
||||||
|
self.embed_error = embed_error
|
||||||
|
self.stream_error = stream_error
|
||||||
|
self.fail_mid_stream = fail_mid_stream
|
||||||
|
self.question_embeds: list[str] = []
|
||||||
|
self.seen_messages: list[list[dict[str, str]]] = []
|
||||||
|
|
||||||
|
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||||
|
self.embed_batches += 1
|
||||||
|
return [_token_vec(t) for t in texts]
|
||||||
|
|
||||||
|
async def embed_one(self, text: str) -> list[float]:
|
||||||
|
if self.embed_error is not None:
|
||||||
|
raise self.embed_error
|
||||||
|
self.question_embeds.append(text)
|
||||||
|
return _token_vec(text)
|
||||||
|
|
||||||
|
async def chat_stream(self, messages: list[dict[str, str]]):
|
||||||
|
self.seen_messages.append(messages)
|
||||||
|
if self.stream_error is not None:
|
||||||
|
raise self.stream_error
|
||||||
|
if self.fail_mid_stream:
|
||||||
|
yield "partial "
|
||||||
|
raise LLMError("mid-stream dropout")
|
||||||
|
for i in range(0, len(self.answer), 12):
|
||||||
|
yield self.answer[i : i + 12]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def seeded_kb(db) -> Iterator[FakeRagLLM]:
|
||||||
|
"""Fresh Postgres with the fixture docs imported (real pipeline)."""
|
||||||
|
db.execute(text("TRUNCATE chunks, documents, query_log"))
|
||||||
|
db.commit()
|
||||||
|
llm = FakeRagLLM()
|
||||||
|
summary = asyncio.run(import_sources([FIXTURES], llm, session=db))
|
||||||
|
assert summary.added == 3
|
||||||
|
yield llm
|
||||||
|
db.execute(text("TRUNCATE chunks, documents, query_log"))
|
||||||
|
db.commit()
|
||||||
|
|
||||||
|
|
||||||
|
def _stream_chat(client: TestClient, message: str) -> tuple[int, str, list[dict[str, Any]]]:
|
||||||
|
with client.stream("POST", "/api/chat", json={"message": message}) as r:
|
||||||
|
assert r.status_code == 200
|
||||||
|
assert r.headers["content-type"].startswith("text/event-stream")
|
||||||
|
buf = ""
|
||||||
|
frames: list[dict[str, Any]] = []
|
||||||
|
for part in r.iter_text():
|
||||||
|
buf += part
|
||||||
|
while "\n\n" in buf:
|
||||||
|
frame, buf = buf.split("\n\n", 1)
|
||||||
|
frame = frame.strip()
|
||||||
|
if frame.startswith("data:"):
|
||||||
|
frames.append(json.loads(frame.removeprefix("data:").strip()))
|
||||||
|
assert buf.strip() == "", "stream must end on a frame boundary"
|
||||||
|
return r.status_code, r.headers["content-type"], frames
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_streams_deltas_then_done_with_sources(client, db, seeded_kb: FakeRagLLM) -> None:
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
|
||||||
|
try:
|
||||||
|
_, _, frames = _stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
deltas = [f for f in frames if f.get("type") == "delta"]
|
||||||
|
assert len(deltas) >= 2 # genuinely streamed
|
||||||
|
assert "".join(d["text"] for d in deltas) == seeded_kb.answer
|
||||||
|
assert not any(f.get("type") == "error" for f in frames)
|
||||||
|
|
||||||
|
done = [f for f in frames if f.get("type") == "done"]
|
||||||
|
assert len(done) == 1
|
||||||
|
assert frames[-1]["type"] == "done" # done is the final event
|
||||||
|
assert done[0]["deflected"] is False
|
||||||
|
assert done[0]["suggestions"] == []
|
||||||
|
sources = done[0]["sources"]
|
||||||
|
assert sources, "done must carry the cited sources"
|
||||||
|
assert sources[0]["path"] == "homelab/kubernetes.md"
|
||||||
|
assert sources[0]["source"] == "docs"
|
||||||
|
assert sources[0]["title"] == "Kubernetes Homelab Cluster"
|
||||||
|
|
||||||
|
# The LLM received the locked HIGH prompt with the FULL document text.
|
||||||
|
(system, user) = seeded_kb.seen_messages[0][0], seeded_kb.seen_messages[0][1]
|
||||||
|
assert user["content"] == QUESTION
|
||||||
|
assert "<relevance>HIGH</relevance>" in system["content"]
|
||||||
|
assert "DEFLECT_MODE" not in system["content"]
|
||||||
|
assert "<documents>" in system["content"]
|
||||||
|
assert "Talos Linux" in system["content"] # full doc, not just the chunk
|
||||||
|
assert "HONESTY GATE" in system["content"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_writes_query_log_row(client, db, seeded_kb: FakeRagLLM) -> None:
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
|
||||||
|
try:
|
||||||
|
_stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
rows = db.scalars(select(QueryLog)).all()
|
||||||
|
assert len(rows) == 1
|
||||||
|
row = rows[0]
|
||||||
|
assert row.question == QUESTION
|
||||||
|
assert row.deflected is False
|
||||||
|
total_chunks = db.scalar(select(func.count()).select_from(Chunk))
|
||||||
|
assert row.chunk_hits == min(get_settings().top_k_chunks, total_chunks)
|
||||||
|
assert row.top_score > 0.0 # genuine token-overlap cosine, best hit
|
||||||
|
assert row.top_score <= 1.0
|
||||||
|
assert "docs/homelab/kubernetes.md" in row.sources
|
||||||
|
assert row.latency_ms >= 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_empty_kb_streams_empty_sources(client, db) -> None:
|
||||||
|
db.execute(text("TRUNCATE chunks, documents, query_log"))
|
||||||
|
db.commit()
|
||||||
|
llm = FakeRagLLM()
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: llm
|
||||||
|
try:
|
||||||
|
_, _, frames = _stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
done = frames[-1]
|
||||||
|
assert done["type"] == "done"
|
||||||
|
assert done["deflected"] is False
|
||||||
|
assert done["sources"] == []
|
||||||
|
row = db.scalars(select(QueryLog)).one()
|
||||||
|
assert row.top_score == 0.0
|
||||||
|
assert row.chunk_hits == 0
|
||||||
|
assert row.sources == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_embed_failure_yields_error_event(client, db, seeded_kb: FakeRagLLM) -> None:
|
||||||
|
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
|
||||||
|
try:
|
||||||
|
_, _, frames = _stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
assert len(frames) == 1
|
||||||
|
assert frames[0]["type"] == "error"
|
||||||
|
assert "embedding" in frames[0]["detail"]
|
||||||
|
assert db.scalars(select(QueryLog)).all() == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_mid_stream_failure_yields_error_after_partial_deltas(client, db) -> None:
|
||||||
|
broken = FakeRagLLM(fail_mid_stream=True)
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
|
||||||
|
try:
|
||||||
|
_, _, frames = _stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
assert [f["type"] for f in frames] == ["delta", "error"]
|
||||||
|
assert "dropped the connection" in frames[1]["detail"]
|
||||||
|
# No done event, no log row for a turn that never completed.
|
||||||
|
assert db.scalars(select(QueryLog)).all() == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_db_down_returns_503_json(client, monkeypatch) -> None:
|
||||||
|
monkeypatch.setattr(chat_api, "db_available", lambda: False)
|
||||||
|
r = client.post("/api/chat", json={"message": "hello"})
|
||||||
|
assert r.status_code == 503
|
||||||
|
assert "offline" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_retrieval_failure_yields_error_event(client, db, seeded_kb, monkeypatch) -> None:
|
||||||
|
def boom(*_a: Any, **_k: Any) -> Any:
|
||||||
|
raise RuntimeError("db exploded")
|
||||||
|
|
||||||
|
monkeypatch.setattr(chat_api, "retrieve", boom)
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
|
||||||
|
try:
|
||||||
|
_, _, frames = _stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
assert [f["type"] for f in frames] == ["error"]
|
||||||
|
assert "offline mid-question" in frames[0]["detail"]
|
||||||
|
assert db.scalars(select(QueryLog)).all() == []
|
||||||
|
|
||||||
|
|
||||||
|
class _BrokenCommitSession:
|
||||||
|
"""Pass-through session whose ``commit()`` raises (query_log failure)."""
|
||||||
|
|
||||||
|
def __init__(self, real: Any) -> None:
|
||||||
|
self._real = real
|
||||||
|
|
||||||
|
def commit(self) -> None:
|
||||||
|
raise RuntimeError("query_log commit failed")
|
||||||
|
|
||||||
|
def __getattr__(self, name: str) -> Any:
|
||||||
|
return getattr(self._real, name)
|
||||||
|
|
||||||
|
|
||||||
|
def test_chat_query_log_failure_still_sends_done(client, db, seeded_kb: FakeRagLLM) -> None:
|
||||||
|
from app.db import SessionLocal
|
||||||
|
|
||||||
|
def broken_db():
|
||||||
|
real = SessionLocal()
|
||||||
|
try:
|
||||||
|
yield _BrokenCommitSession(real)
|
||||||
|
finally:
|
||||||
|
real.close()
|
||||||
|
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_db] = broken_db
|
||||||
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
|
||||||
|
try:
|
||||||
|
_, _, frames = _stream_chat(client, QUESTION)
|
||||||
|
finally:
|
||||||
|
fastapi_app.dependency_overrides.clear()
|
||||||
|
|
||||||
|
# The answer (and the done event) went out despite the log-row failure.
|
||||||
|
assert [f["type"] for f in frames if f["type"] == "delta"]
|
||||||
|
assert frames[-1]["type"] == "done"
|
||||||
|
assert frames[-1]["deflected"] is False
|
||||||
@@ -10,12 +10,13 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
|
from types import SimpleNamespace
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from app.config import Settings
|
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:
|
def _settings(**kwargs: Any) -> Settings:
|
||||||
@@ -230,3 +231,98 @@ def test_single_oversized_text_fails_actionably() -> None:
|
|||||||
with pytest.raises(EmbeddingError, match="token cap"):
|
with pytest.raises(EmbeddingError, match="token cap"):
|
||||||
asyncio.run(llm.embed(["x" * 3000]))
|
asyncio.run(llm.embed(["x" * 3000]))
|
||||||
assert llm.embed_batches == 0
|
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