324 lines
12 KiB
Python
324 lines
12 KiB
Python
"""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?"
|
|
OFF_TOPIC = "How do I bake sourdough bread?"
|
|
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 >= get_settings().relevance_threshold # why the gate answered
|
|
assert row.top_score <= 1.0
|
|
assert "docs/homelab/kubernetes.md" in row.sources
|
|
assert row.latency_ms >= 0
|
|
|
|
|
|
def test_off_topic_question_deflects_honestly(client, db, seeded_kb: FakeRagLLM) -> None:
|
|
"""Phase 04 contract: weak retrieval ⇒ honest deflection, no fake answer."""
|
|
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: seeded_kb
|
|
try:
|
|
_, _, frames = _stream_chat(client, OFF_TOPIC)
|
|
finally:
|
|
fastapi_app.dependency_overrides.clear()
|
|
|
|
assert not any(f.get("type") == "error" for f in frames)
|
|
deltas = [f for f in frames if f.get("type") == "delta"]
|
|
assert len(deltas) >= 2 # the LLM is still called (voice stays chippy)
|
|
done = frames[-1]
|
|
assert done["type"] == "done"
|
|
assert done["deflected"] is True
|
|
# 2-3 alternative chips, all non-empty, derived from real titles/topics.
|
|
assert 2 <= len(done["suggestions"]) <= 3
|
|
assert all(s.strip() for s in done["suggestions"])
|
|
assert any(
|
|
"Deploying a New Service" in s for s in done["suggestions"]
|
|
), "the best weak-hit title must be offered as a chip"
|
|
assert done["sources"], "weak hits are still reported as the closest sources"
|
|
|
|
# The LLM saw the LOW prompt: DEFLECT_MODE + titles, never doc content.
|
|
(system, user) = seeded_kb.seen_messages[0][0], seeded_kb.seen_messages[0][1]
|
|
assert user["content"] == OFF_TOPIC
|
|
assert "<relevance>LOW</relevance>" in system["content"]
|
|
assert "DEFLECT_MODE" in system["content"]
|
|
assert "HONESTY GATE" in system["content"]
|
|
assert "Talos Linux" not in system["content"] # full doc content never sent
|
|
assert "<documents>" not in system["content"]
|
|
|
|
# Durable record: deflected=true + the weak top_score.
|
|
row = db.scalars(select(QueryLog)).one()
|
|
assert row.question == OFF_TOPIC
|
|
assert row.deflected is True
|
|
assert 0.0 < row.top_score < get_settings().relevance_threshold
|
|
assert row.chunk_hits >= 1
|
|
|
|
|
|
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()
|
|
|
|
# Nothing retrieved ⇒ nothing to pretend to know: honest deflection.
|
|
done = frames[-1]
|
|
assert done["type"] == "done"
|
|
assert done["deflected"] is True
|
|
assert done["sources"] == []
|
|
assert 2 <= len(done["suggestions"]) <= 3 # onboarding fallback chips
|
|
(system, _user) = llm.seen_messages[0][0], llm.seen_messages[0][1]
|
|
assert "DEFLECT_MODE" in system["content"]
|
|
assert "nothing close at all" in system["content"]
|
|
row = db.scalars(select(QueryLog)).one()
|
|
assert row.deflected is True
|
|
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
|