462 lines
20 KiB
Python
462 lines
20 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).
|
||
|
||
Summaries (phase 30): a lite-model summary chunk's parent *is* the
|
||
source document, so a summary hit resolves to the full source document
|
||
through the unchanged chunk→document mapping (A7 revised) — context
|
||
assembly is untouched. ``TurnPlan.summary_hits`` counts the hit chunks
|
||
with ``is_summary`` whose parent document landed in the selected
|
||
top-N context, and the per-turn log line records ``summary_hits=N``
|
||
after ``fts_hits`` (PLAN §9 line extension).
|
||
|
||
KB overview (phase 31): the lite-generated outline of the knowledge
|
||
base (single ``kb_overview`` row) is read per turn (one indexed PK
|
||
lookup — no LLM call) and injected into **both** prompts as the
|
||
``<knowledge_base>`` section, ordered ``<relevance>`` →
|
||
``<knowledge_base>`` → ``<tuning>`` → mode body. With an empty row the
|
||
prompts stay byte-identical to the pre-phase text (phase 15
|
||
convention); ``TurnPlan.kb_chars`` records the length of the stored
|
||
outline (0 when absent) and the per-turn log line records
|
||
``kb_chars=N`` after ``tuning=N`` (PLAN §9 line extension).
|
||
|
||
Agent document tools (phase 37, PLAN §4 extension, owner permission
|
||
2026-08-26; phase 45 removed the per-tool budgets — owner permission
|
||
2026-08-27): a **grounded** turn (``not plan.deflected``) no longer
|
||
streams a bare ``chat_stream`` — it runs the agent loop
|
||
(``app.rag.agent.run_agent``), which offers the model the two
|
||
server-side tools ``list_documents`` / ``read_document`` for the whole
|
||
turn (as many calls as the model wants, re-lists included) until it
|
||
answers or the round cap (``BOR_AGENT_MAX_ROUNDS``, default 10) forces
|
||
one final no-tools answer. Each model-requested call streams as an SSE
|
||
``tool`` event — ``{"type": "tool", "name": …, "argument":
|
||
"source/path" | null}`` — ahead of the answer's ``delta`` frames.
|
||
``done.sources``, ``query_log.sources`` and the per-turn log line all
|
||
report the same combined source list (retrieval docs + the agent's
|
||
read docs, deduped by ``(source, path)``, order preserved), and the log
|
||
line records ``tool_calls=N`` after ``thinking_chars=N`` (PLAN §9 line
|
||
extension — ``N`` counts executed tool calls; rejected calls do not
|
||
count). **Deflected turns keep the direct ``chat_stream`` —
|
||
byte-identical to the pre-phase path (A8):** the LOW prompt never
|
||
carries tools, and with ``agent_max_rounds`` at **0** ``run_agent``
|
||
makes exactly one ``tools=None`` request, reproducing the pre-phase
|
||
behavior (the kill switch).
|
||
"""
|
||
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.agent import AgentHolder, run_agent
|
||
from app.rag.llm import (
|
||
EmbeddingError,
|
||
LLMClient,
|
||
LLMError,
|
||
StreamPiece, # type of the answer pieces streamed by the agent loop
|
||
ToolCallPiece, # phase 37: one model-requested tool call
|
||
)
|
||
from app.rag.overview import load_kb_overview
|
||
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,
|
||
ChatToolEvent,
|
||
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
|
||
#: Hit chunks with ``is_summary`` whose parent document made it into
|
||
#: *docs* (phase 30; per-turn log line ``summary_hits=N``).
|
||
summary_hits: int = 0
|
||
#: Length of the stored KB overview injected as the
|
||
#: ``<knowledge_base>`` section (phase 31; per-turn log line
|
||
#: ``kb_chars=N``). 0 when no non-empty row exists.
|
||
kb_chars: int = 0
|
||
|
||
|
||
def plan_turn(
|
||
chunks: Sequence[RetrievedChunk],
|
||
settings: Settings,
|
||
notes: Sequence[str] | None = None,
|
||
kb_overview: 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.
|
||
|
||
*kb_overview* is the stored KB outline (phase 31, one PK lookup per
|
||
turn): when non-empty, both prompts carry the ``<knowledge_base>``
|
||
section (between ``<relevance>`` and ``<tuning>``) and
|
||
``kb_chars`` records the outline's length; with no outline the
|
||
prompts are byte-identical to the pre-phase text and ``kb_chars``
|
||
is 0.
|
||
|
||
``summary_hits`` (phase 30) counts the hit chunks with
|
||
``is_summary`` whose parent document is among the selected
|
||
top-N documents — both the HIGH and the LOW branch record it.
|
||
"""
|
||
steering = list(notes or [])
|
||
kb_text = (kb_overview or "").strip()
|
||
kb_chars = len(kb_text)
|
||
best_cosine = max((c.cosine for c in chunks), default=0.0)
|
||
fts_hits = sum(1 for c in chunks if c.fts_hit)
|
||
docs = select_documents(chunks, n=settings.top_n_docs)
|
||
selected_ids = {d.id for d in docs}
|
||
summary_hits = sum(1 for c in chunks if c.is_summary and c.document.id in selected_ids)
|
||
if best_cosine >= settings.relevance_threshold or fts_hits > 0:
|
||
return TurnPlan(
|
||
best_cosine,
|
||
fts_hits,
|
||
False,
|
||
build_high_prompt(docs, notes=steering, kb_overview=kb_text),
|
||
docs,
|
||
[],
|
||
len(steering),
|
||
summary_hits,
|
||
kb_chars,
|
||
)
|
||
titles = weak_hit_titles(chunks)
|
||
return TurnPlan(
|
||
best_cosine,
|
||
fts_hits,
|
||
True,
|
||
build_deflect_prompt(titles, notes=steering, kb_overview=kb_text),
|
||
docs,
|
||
derive_suggestions(titles, settings.suggestions),
|
||
len(steering),
|
||
summary_hits,
|
||
kb_chars,
|
||
)
|
||
|
||
|
||
@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]:
|
||
# Phase 48: one terminal flag — ``True`` at every terminal exit
|
||
# (the ``done`` yield; every ``error``-then-``return``). The
|
||
# ``finally`` below logs the cancelled-turn line only when the
|
||
# consumer went away before any terminal frame; it must not
|
||
# yield (GeneratorExit handling).
|
||
settled = False
|
||
try:
|
||
# 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,
|
||
)
|
||
settled = True # terminal: the error frame settles the turn
|
||
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)
|
||
# KB overview (phase 31): one indexed PK lookup per turn —
|
||
# the outline is generated at import time, never per chat
|
||
# turn.
|
||
kb_overview = load_kb_overview(db)
|
||
chunks = retrieve(db, request.message, question_vec)
|
||
plan = plan_turn(chunks, settings, notes=steering_notes, kb_overview=kb_overview)
|
||
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),
|
||
)
|
||
settled = True # terminal: the error frame settles the turn
|
||
yield sse_event(
|
||
ChatErrorEvent(
|
||
detail="The knowledge base went offline mid-question — is Postgres up?"
|
||
).model_dump()
|
||
)
|
||
return
|
||
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.
|
||
# Phase 37: a grounded turn runs the agent loop instead of
|
||
# a bare ``chat_stream`` — its ``ToolCallPiece``s stream
|
||
# as ``tool`` events ahead of the answer. A deflected turn
|
||
# keeps the direct ``chat_stream`` (byte-identical, A8):
|
||
# the LOW prompt never carries tools, and with
|
||
# ``agent_max_rounds=0`` ``run_agent`` is a single
|
||
# ``tools=None`` request anyway (the kill switch).
|
||
holder = AgentHolder()
|
||
answer_stream: AsyncIterator[StreamPiece | ToolCallPiece]
|
||
if plan.deflected:
|
||
answer_stream = llm.chat_stream(messages)
|
||
else:
|
||
answer_stream = run_agent(
|
||
llm,
|
||
db,
|
||
system_prompt=plan.system_prompt,
|
||
user_message=request.message,
|
||
seed_docs=plan.docs,
|
||
settings=settings,
|
||
holder=holder,
|
||
)
|
||
thinking_chars = 0
|
||
try:
|
||
async for piece in answer_stream: # StreamPiece | ToolCallPiece
|
||
if isinstance(piece, ToolCallPiece):
|
||
# Phase 37 (PLAN §4 extension): one SSE ``tool``
|
||
# frame per model-requested call; ``argument`` is
|
||
# the read_document "source/path" (null
|
||
# otherwise).
|
||
yield sse_event(
|
||
ChatToolEvent(
|
||
name=piece.name,
|
||
argument=(
|
||
f"{piece.arguments.get('source')}/{piece.arguments.get('path')}"
|
||
if piece.name == "read_document"
|
||
else None
|
||
),
|
||
).model_dump()
|
||
)
|
||
continue
|
||
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,
|
||
)
|
||
settled = True # terminal: the error frame settles the turn
|
||
yield sse_event(
|
||
ChatErrorEvent(
|
||
detail="The chat model dropped the connection — try again?"
|
||
).model_dump()
|
||
)
|
||
return
|
||
except Exception: # noqa: BLE001 — a tool call hit the DB mid-stream
|
||
# Phase 37: tool execution (list_catalog / find_document)
|
||
# runs inside the stream now; a mid-turn DB failure gets
|
||
# the same structured ``error`` event as the pre-stream
|
||
# retrieval path.
|
||
logger.exception(
|
||
"chat: tool execution failed question=%r total_ms=%d",
|
||
request.message,
|
||
int((time.monotonic() - started) * 1000),
|
||
)
|
||
settled = True # terminal: the error frame settles the turn
|
||
yield sse_event(
|
||
ChatErrorEvent(
|
||
detail="The knowledge base went offline mid-question — is Postgres up?"
|
||
).model_dump()
|
||
)
|
||
return
|
||
|
||
# 4. Durable record + required per-turn log line (PLAN §9).
|
||
# Phase 37: the agent's read documents join the
|
||
# retrieval's — deduped by (source, path), order preserved
|
||
# — and the same combined list feeds done.sources,
|
||
# query_log.sources and the log line (empty on deflected
|
||
# turns: the agent never runs). A cancelled turn (the
|
||
# generator closed by the consumer) never reaches this
|
||
# step — no query_log row.
|
||
cited_docs: list[Document] = []
|
||
seen: set[tuple[str, str]] = set()
|
||
for doc in [*plan.docs, *holder.read_docs]:
|
||
key = (doc.source, doc.path)
|
||
if key not in seen:
|
||
seen.add(key)
|
||
cited_docs.append(doc)
|
||
source_paths = [f"{d.source}/{d.path}" for d in cited_docs]
|
||
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 summary_hits=%d tuning=%d "
|
||
"kb_chars=%d threshold=%.2f deflected=%s sources=%r thinking_chars=%d "
|
||
"tool_calls=%d total_ms=%d",
|
||
request.message,
|
||
embed_ms,
|
||
plan.top_score,
|
||
plan.fts_hits,
|
||
plan.summary_hits,
|
||
plan.tuning_count,
|
||
plan.kb_chars,
|
||
settings.relevance_threshold,
|
||
plan.deflected,
|
||
source_paths,
|
||
thinking_chars,
|
||
holder.tool_calls,
|
||
total_ms,
|
||
)
|
||
settled = True # terminal: the done frame settles the turn
|
||
yield sse_event(
|
||
ChatDoneEvent(
|
||
deflected=plan.deflected,
|
||
sources=[
|
||
SourceRef(source=d.source, path=d.path, title=d.title) for d in cited_docs
|
||
],
|
||
suggestions=plan.suggestions,
|
||
).model_dump()
|
||
)
|
||
finally:
|
||
# Phase 48 (owner-locked): a cancelled turn — the SSE
|
||
# consumer went away before any terminal frame — settles
|
||
# with one warning line and skips query_log entirely (the
|
||
# write above is simply never reached when the generator is
|
||
# closed). The finally must not yield (GeneratorExit
|
||
# handling).
|
||
if not settled:
|
||
logger.warning(
|
||
"chat: turn cancelled question=%r total_ms=%d",
|
||
request.message,
|
||
int((time.monotonic() - started) * 1000),
|
||
)
|
||
|
||
return StreamingResponse(stream(), media_type="text/event-stream", headers=SSE_HEADERS)
|