291 lines
11 KiB
Python
291 lines
11 KiB
Python
"""POST /api/chat — a RAG chat turn streamed over SSE (PLAN §3/§4).
|
||
|
||
Flow (LOCKED A7/A15): embed the question → hybrid retrieval (cosine
|
||
top-N ∪ Postgres FTS top-N, RRF-fused) → the **honesty gate** → 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.
|
||
|
||
Thinking (phase 17, PLAN §4 extension, owner permission 2026-08-23): the
|
||
model's reasoning arrives ahead of the answer and is streamed as
|
||
``thinking`` events before the ``delta`` events of the same turn. Each
|
||
turn's thinking is counted in the per-turn log line
|
||
(``thinking_chars=N``); ``BOR_STREAM_THINKING=0`` suppresses the
|
||
``thinking`` frames (the pieces are still counted).
|
||
|
||
Honesty gate (A8, revised 2026-08-21): LOW — deflection — only when the
|
||
best cosine is strictly below ``BOR_RELEVANCE_THRESHOLD`` **and** no
|
||
candidate chunk FTS-matches the question (``fts_hits == 0``). A
|
||
name-your-tool question with weak vector overlap but a lexical hit still
|
||
gets a grounded answer. Deflection mode carries weak-hit *titles only*
|
||
(never document content) plus deterministic "Maybe try" chips, and the
|
||
``done`` event / ``query_log`` row record ``deflected=true``, the weak
|
||
score and the ``fts_hits`` count.
|
||
|
||
Steering (phase 15): the owner's stored tuning notes are loaded per turn
|
||
(oldest first) and injected into the system prompt as a ``<tuning>``
|
||
section — both the HIGH and the LOW prompt carry it. The per-turn log
|
||
line records ``tuning=N`` (the number of injected notes).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import time
|
||
from collections.abc import AsyncIterator, Sequence
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from fastapi import APIRouter, Depends
|
||
from fastapi.responses import JSONResponse, StreamingResponse
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.api.steering import load_steering_notes
|
||
from app.config import Settings, get_settings
|
||
from app.db import db_available, get_db
|
||
from app.models import Document, QueryLog
|
||
from app.rag.llm import (
|
||
EmbeddingError,
|
||
LLMClient,
|
||
LLMError,
|
||
StreamPiece, # noqa: F401 (phase 17 typing: chat_stream yields StreamPiece)
|
||
)
|
||
from app.rag.prompts import build_deflect_prompt, build_high_prompt
|
||
from app.rag.retriever import RetrievedChunk, retrieve, select_documents, weak_hit_titles
|
||
from app.rag.suggestions import derive_suggestions
|
||
from app.schemas import (
|
||
ChatDoneEvent,
|
||
ChatErrorEvent,
|
||
ChatRequest,
|
||
ChatThinkingEvent,
|
||
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"
|
||
|
||
|
||
@dataclass
|
||
class TurnPlan:
|
||
"""What one chat turn sends to the LLM and reports on ``done``."""
|
||
|
||
top_score: float # best cosine across candidates (query_log.top_score)
|
||
fts_hits: int # lexical (OR-tsquery) candidates matched
|
||
deflected: bool
|
||
system_prompt: str
|
||
docs: list[Document] # cited sources (weak hits when deflected)
|
||
suggestions: list[str] # "Maybe try" chips (deflected turns only)
|
||
tuning_count: int = 0 # steering notes injected into the system prompt
|
||
|
||
|
||
def plan_turn(
|
||
chunks: Sequence[RetrievedChunk],
|
||
settings: Settings,
|
||
notes: Sequence[str] | None = None,
|
||
) -> TurnPlan:
|
||
"""Apply the honesty gate (A8, revised) and assemble prompt + context.
|
||
|
||
* **HIGH (grounded)** when ``best_cosine >= threshold`` **or**
|
||
``fts_hits > 0``: HIGH prompt with the full top-N documents, no
|
||
suggestions. A cosine exactly at the threshold is an answer — the
|
||
gate is strict (``< threshold``).
|
||
* **LOW (deflected)** only when ``best_cosine < threshold`` **and**
|
||
``fts_hits == 0`` (or no hits at all): LOW prompt (``DEFLECT_MODE``)
|
||
with weak-hit titles only — never document content — plus
|
||
deterministic alternative-question chips derived from those titles.
|
||
|
||
``top_score`` (stored in ``query_log``) is the best cosine, so the
|
||
gate input is always a pure vector-similarity number; the lexical
|
||
signal is recorded separately as ``fts_hits``.
|
||
|
||
*notes* are the owner's steering notes (phase 15, oldest first):
|
||
when non-empty, both the HIGH and the LOW prompt carry the
|
||
``<tuning>`` section; with no notes the prompts are unchanged.
|
||
"""
|
||
steering = list(notes or [])
|
||
best_cosine = max((c.cosine for c in chunks), default=0.0)
|
||
fts_hits = sum(1 for c in chunks if c.fts_hit)
|
||
if best_cosine >= settings.relevance_threshold or fts_hits > 0:
|
||
docs = select_documents(chunks, n=settings.top_n_docs)
|
||
return TurnPlan(
|
||
best_cosine,
|
||
fts_hits,
|
||
False,
|
||
build_high_prompt(docs, notes=steering),
|
||
docs,
|
||
[],
|
||
len(steering),
|
||
)
|
||
titles = weak_hit_titles(chunks)
|
||
return TurnPlan(
|
||
best_cosine,
|
||
fts_hits,
|
||
True,
|
||
build_deflect_prompt(titles, notes=steering),
|
||
select_documents(chunks, n=settings.top_n_docs),
|
||
derive_suggestions(titles, settings.suggestions),
|
||
len(steering),
|
||
)
|
||
|
||
|
||
@router.post("/chat")
|
||
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:
|
||
embed_ms = int((time.monotonic() - t0) * 1000)
|
||
total_ms = int((time.monotonic() - started) * 1000)
|
||
logger.error(
|
||
"chat: question=%r embed_ms=%d total_ms=%d — embedding failed: %s",
|
||
request.message,
|
||
embed_ms,
|
||
total_ms,
|
||
e,
|
||
)
|
||
yield sse_event(
|
||
ChatErrorEvent(
|
||
detail="I couldn't reach the embedding model — please try again."
|
||
).model_dump()
|
||
)
|
||
return
|
||
embed_ms = int((time.monotonic() - t0) * 1000)
|
||
|
||
# 2. Retrieve top-K chunks, load the owner's steering notes
|
||
# (phase 15), then the honesty gate (A8) picks the HIGH
|
||
# (grounded) or LOW (deflected) prompt + context.
|
||
settings = get_settings()
|
||
try:
|
||
steering_notes = load_steering_notes(db)
|
||
chunks = retrieve(db, request.message, question_vec)
|
||
plan = plan_turn(chunks, settings, notes=steering_notes)
|
||
except Exception: # noqa: BLE001 — DB failure mid-turn
|
||
logger.exception(
|
||
"chat: retrieval failed question=%r total_ms=%d",
|
||
request.message,
|
||
int((time.monotonic() - started) * 1000),
|
||
)
|
||
yield sse_event(
|
||
ChatErrorEvent(
|
||
detail="The knowledge base went offline mid-question — is Postgres up?"
|
||
).model_dump()
|
||
)
|
||
return
|
||
source_paths = [f"{d.source}/{d.path}" for d in plan.docs]
|
||
messages = [
|
||
{"role": "system", "content": plan.system_prompt},
|
||
{"role": "user", "content": request.message},
|
||
]
|
||
|
||
# 3. Stream the answer (grounded, or an honest deflection).
|
||
# Phase 17: thinking pieces stream as ``thinking`` events
|
||
# ahead of the ``delta`` events (PLAN §4 extension); the
|
||
# kill-switch (``BOR_STREAM_THINKING=0``) suppresses the
|
||
# frames, not the counting.
|
||
thinking_chars = 0
|
||
try:
|
||
async for piece in llm.chat_stream(messages): # StreamPiece (phase 17)
|
||
if piece.kind == "thinking":
|
||
thinking_chars += len(piece.text)
|
||
if settings.stream_thinking:
|
||
yield sse_event(ChatThinkingEvent(text=piece.text).model_dump())
|
||
else:
|
||
yield sse_event({"type": "delta", "text": piece.text})
|
||
except LLMError as e:
|
||
logger.error(
|
||
"chat: LLM stream failed question=%r total_ms=%d — %s",
|
||
request.message,
|
||
int((time.monotonic() - started) * 1000),
|
||
e,
|
||
)
|
||
yield sse_event(
|
||
ChatErrorEvent(
|
||
detail="The chat model dropped the connection — try again?"
|
||
).model_dump()
|
||
)
|
||
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=plan.top_score,
|
||
fts_hits=plan.fts_hits,
|
||
chunk_hits=len(chunks),
|
||
deflected=plan.deflected,
|
||
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 fts_hits=%d tuning=%d threshold=%.2f "
|
||
"deflected=%s sources=%r thinking_chars=%d total_ms=%d",
|
||
request.message,
|
||
embed_ms,
|
||
plan.top_score,
|
||
plan.fts_hits,
|
||
plan.tuning_count,
|
||
settings.relevance_threshold,
|
||
plan.deflected,
|
||
source_paths,
|
||
thinking_chars,
|
||
total_ms,
|
||
)
|
||
yield sse_event(
|
||
ChatDoneEvent(
|
||
deflected=plan.deflected,
|
||
sources=[
|
||
SourceRef(source=d.source, path=d.path, title=d.title) for d in plan.docs
|
||
],
|
||
suggestions=plan.suggestions,
|
||
).model_dump()
|
||
)
|
||
|
||
return StreamingResponse(stream(), media_type="text/event-stream", headers=SSE_HEADERS)
|