diff --git a/app/api/chat.py b/app/api/chat.py index 67842b1..c05f019 100644 --- a/app/api/chat.py +++ b/app/api/chat.py @@ -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: \\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) diff --git a/app/rag/llm.py b/app/rag/llm.py index dba12c2..1d60688 100644 --- a/app/rag/llm.py +++ b/app/rag/llm.py @@ -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 diff --git a/app/rag/prompts.py b/app/rag/prompts.py new file mode 100644 index 0000000..9ec5b65 --- /dev/null +++ b/app/rag/prompts.py @@ -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 ````. +* ``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 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}" +) + + +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'\n' + f"{doc.content}\n" + "" + for doc in documents + ] + body = "\n\n".join(blocks) if blocks else ( + "(no documents matched — do not invent specifics)" + ) + return _base("HIGH") + "\n\n" + body + "\n" + + +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 + ) diff --git a/app/rag/retriever.py b/app/rag/retriever.py new file mode 100644 index 0000000..f137d30 --- /dev/null +++ b/app/rag/retriever.py @@ -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 diff --git a/frontend/assets/app.js b/frontend/assets/app.js index 062f12d..ae403ee 100644 --- a/frontend/assets/app.js +++ b/frontend/assets/app.js @@ -1,8 +1,11 @@ /* Brain of Reese — chat shell. * - * Scaffolding-stage behavior: renders suggestions, shows KB health, and - * echoes a friendly placeholder answer. The real RAG streaming chat is - * implemented in the chat-rag phase (see .agent/user_stories/). + * Renders suggestions, shows KB health, and runs chat turns against + * POST /api/chat (SSE, PLAN §4): deltas render live into the Brain bubble, + * 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. */ @@ -148,6 +151,70 @@ function autoGrow() { 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) { e.preventDefault(); const text = input.value.trim(); @@ -156,21 +223,56 @@ async function handleSend(e) { addMessage("user", renderMarkdown(text)); input.value = ""; autoGrow(); + clearErrorBanner(); setBusy(true); addTyping(); + let wrap = null; + let acc = ""; + let res = null; try { - // TODO(chat-rag phase): replace with POST /api/chat (SSE streaming). - await new Promise((res) => setTimeout(res, 500)); - const reply = - "I'm still getting my neurons wired up — the real me ships in the " + - "next phase! Keep the questions coming, you're on a roll. 🚀"; + res = await fetch("/api/chat", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ message: text }), + }); + 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(); - addMessage("brain", renderMarkdown(reply)); - } catch { - removeTyping(); - addMessage("brain", "Something went wrong on my side — please try again in a moment."); + showErrorBanner(err.message || "Something went wrong on my side."); } finally { + try { res?.body?.cancel(); } catch { /* stream already closed */ } setBusy(false); input.focus(); } diff --git a/frontend/assets/styles.css b/frontend/assets/styles.css index 3876f8d..7cf7708 100644 --- a/frontend/assets/styles.css +++ b/frontend/assets/styles.css @@ -56,6 +56,11 @@ body { } /* ---------- 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 { position: absolute !important; width: 1px; height: 1px; diff --git a/tests/e2e/test_chat_rag.py b/tests/e2e/test_chat_rag.py new file mode 100644 index 0000000..3172486 --- /dev/null +++ b/tests/e2e/test_chat_rag.py @@ -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"]) diff --git a/tests/e2e/test_smoke.py b/tests/e2e/test_smoke.py index 825d742..af1309d 100644 --- a/tests/e2e/test_smoke.py +++ b/tests/e2e/test_smoke.py @@ -1,12 +1,10 @@ -"""Phase 01 smoke E2E: the app boots, serves the local frontend, and the -placeholder chat round-trips without a stale button. +"""Phase 01 smoke E2E: the app boots, serves the local frontend, and a chat +round-trip never leaves a stale button (answer or error banner, both fine). Run: uv run pytest tests/e2e/test_smoke.py -v """ from __future__ import annotations -import re - import httpx 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', "") -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.locator("#message-input").fill("hello brain") 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) - brain_bubble = page.locator(".msg.brain .bubble").first - brain_bubble.wait_for(state="visible", timeout=10_000) - # to_have_text retries until the async fetch resolves (no stale read). - expect(brain_bubble).to_have_text(re.compile("neurons"), timeout=10_000) + # Then either a streamed Brain answer (DB up) or an error banner + # (DB down) — but the turn must always complete. + page.wait_for_selector( + ".msg.brain .bubble, #kb-banner.is-error", + state="visible", + timeout=20_000, + ) # Button is never left stuck: back to "Send" and enabled. - btn = page.locator("#send-btn") - assert btn.is_enabled() - assert "Send" in btn.inner_text() + expect(page.locator("#send-btn")).to_be_enabled(timeout=10_000) + expect(page.locator("#send-label")).to_have_text("Send", timeout=10_000) diff --git a/tests/integration/test_api.py b/tests/integration/test_api.py index a015652..51eb238 100644 --- a/tests/integration/test_api.py +++ b/tests/integration/test_api.py @@ -34,16 +34,6 @@ def test_styles_and_js_served(client) -> None: 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: r = client.post("/api/chat", json={"message": ""}) assert r.status_code == 422 diff --git a/tests/integration/test_chat_api.py b/tests/integration/test_chat_api.py new file mode 100644 index 0000000..21e5b46 --- /dev/null +++ b/tests/integration/test_chat_api.py @@ -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 "HIGH" in system["content"] + assert "DEFLECT_MODE" not in system["content"] + assert "" 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 diff --git a/tests/unit/test_llm_client.py b/tests/unit/test_llm_client.py index ffeb2c2..a5c68d7 100644 --- a/tests/unit/test_llm_client.py +++ b/tests/unit/test_llm_client.py @@ -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"}])) diff --git a/tests/unit/test_prompts.py b/tests/unit/test_prompts.py new file mode 100644 index 0000000..c8df553 --- /dev/null +++ b/tests/unit/test_prompts.py @@ -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 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 "HIGH" in prompt + assert "DEFLECT_MODE" not in prompt + assert "" in prompt and "" 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 "" 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 "LOW" 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 "" 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") diff --git a/tests/unit/test_retriever.py b/tests/unit/test_retriever.py new file mode 100644 index 0000000..ea3503d --- /dev/null +++ b/tests/unit/test_retriever.py @@ -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) == [] diff --git a/tests/unit/test_sse_events.py b/tests/unit/test_sse_events.py new file mode 100644 index 0000000..23449f8 --- /dev/null +++ b/tests/unit/test_sse_events.py @@ -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"