ducoterra
3a4035fc96
--- **Phase 110 — Fix SSE DB Connection Pool Exhaustion (SEC-14-04): COMPLETE** **What was implemented/verified:** - All three tasks (pool config, short-lived sessions, concurrency cap) were already implemented in code - Fixed `tests/integration/test_chat_db_sessions.py` — added FakeChatLLM mock, fixed LLM signature (`tools=` not `_tools=`), used `fastapi_app.dependency_overrides` instead of `client.app.dependency_overrides` - Fixed `tests/e2e/test_chat_db_pool.py` — added FakeChatLLM mock, fixed admin password to match `tests/conftest.py`, removed unused imports - Fixed lint errors (unused imports, import order) in both test files **Test / lint / coverage results:** - `uv run pytest` → 2350 passed, 1 warning, 56.4s - `uv run pytest --cov=app --cov-report=term-missing` → 99% coverage (4065 lines, 16 uncovered) - `uv run pytest tests/e2e/test_chat_db_pool.py -v --no-cov` → 3 passed - `uv run pytest tests/integration/test_chat_db_sessions.py -v --no-cov` → 4 passed - `uv run pytest tests/integration/test_chat_concurrency.py -v --no-cov` → 11 passed - `uv run pytest tests/unit/test_db_pool_config.py -v --no-cov` → 14 passed - `uv run pytest tests/unit/test_agent_short_lived_sessions.py -v --no-cov` → 7 passed - `uv run ruff check .` → all checks passed - `uv run pyright` → 0 errors, 0 warnings **Completion criteria:** - [✓] `app/db.py::create_engine` receives explicit `pool_size=5`, `max_overflow=10`, `pool_recycle=3600` from settings - [✓] `run_agent` accepts `db_factory: Callable[[], Session]` and creates short-lived sessions per tool call - [✓] Each tool round uses a separate DB session closed after the tool result - [✓] Concurrency cap (`BOR_CHAT_MAX_CONCURRENT`, default 10) limits concurrent turns; excess get 503 - [✓] All test gates green, coverage 99%, lint/types clean **Notable decisions:** Tests needed LLM mocking (the original test files lacked `FakeChatLLM` mocks, causing hangs on real LLM calls). **Next pending phase:** None — this is the last phase in `todo/`.
823 lines
41 KiB
Python
823 lines
41 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; phase 70 aligned the surface to the harness-trained
|
||
``ls`` / ``read`` / ``grep`` — owner permission 2026-09-03): 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 three server-side tools
|
||
``ls`` / ``read`` / ``grep`` for the whole turn (as many calls as the
|
||
model wants, re-lists and re-greps 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": … | null}`` — ahead
|
||
of the answer's ``delta`` frames: ``argument`` is the single string the
|
||
model passed — ``read``'s ``path`` (the combined ``source/path``),
|
||
``grep``'s ``pattern``, ``ls``'s ``path`` — or null (a non-string
|
||
value — a model error the backend refuses — and an omitted argument
|
||
both yield null).
|
||
|
||
Tool-result frames (phase 95, ``TODO.md`` L5 — A15 extension, owner
|
||
permission 2026-09-10; the event-type list grows from six to SEVEN,
|
||
``tool_result`` among them; PLAN.md is being redone by the owner): a
|
||
``read`` whose document is longer than ``BOR_READ_MAX_CHARS`` streams,
|
||
AFTER the matching ``tool`` frame (the line is already on screen — the
|
||
marker lands a beat later, the phase-37/48 tool-line timing is
|
||
untouched) and BEFORE the next model round, exactly one optional
|
||
``tool_result`` event — ``{"type": "tool_result", "name": "read",
|
||
"argument": …, "truncated": true, "chars_shown": N, "chars_total": M}``
|
||
— the additive truncation notice the UI turns into the
|
||
"(truncated — showing N of M chars)" marker on the Reading line. A
|
||
non-truncated read streams NO such frame (one frame = one noteworthy
|
||
event), deflected turns never stream one (the agent never runs, A8),
|
||
and every pre-existing frame is byte-identical — clients that do not
|
||
know the type ignore it. ``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 — a grep adds no source; it is a locator, locked A5), 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).
|
||
|
||
LLM retries (phase 67, ``BOR_LLM_RETRIES`` / ``BOR_LLM_RETRY_DELAY``,
|
||
owner-locked 2026-09-01): when the aipi endpoint dies before a request
|
||
has streamed its first output frame (locked A2), the request is
|
||
restarted — up to ``llm_retries`` times (default 3), a flat
|
||
``llm_retry_delay`` (default 5 s) between attempts. Every restart is
|
||
announced with an SSE ``retry`` frame (``{"type": "retry",
|
||
"attempt": n, "max_attempts": N}`` — the attempt about to be tried,
|
||
1-based, ahead of the pre-retry wait) so the UI can show the transient
|
||
"Communication interrupted — retrying (n of N)…" status (locked A4);
|
||
exhaustion and any failure after the first frame keep the existing
|
||
terminal ``error`` frames. The pre-stream question embedding retries on
|
||
the same budget, and the deflected answer stream goes through
|
||
``chat_stream_retried`` (the agent loop retries per round — task 03).
|
||
The per-turn log line records ``retries=N`` after ``total_ms=N`` (0
|
||
when nothing was retried — the field is uniform across all turn
|
||
shapes).
|
||
|
||
Tool-scaffolding guardrail (phase 71, deterministic only — owner
|
||
permission 2026-09-03: "deterministic guardrails only right now, forget
|
||
using a model for that"): the raw chat-template tokens
|
||
``<|tool_call_start|>…<|tool_call_end|>`` the ``lite`` model sometimes
|
||
emits as plain answer text can never reach the user. The DEFLECTED
|
||
path's request runs ``delta.content`` through a caller-owned
|
||
``ScaffoldingFilter`` (content only — thinking stays raw); when the
|
||
filter wipes the whole reply (visible content 0, ``stripped_chars > 0``)
|
||
the turn gets exactly ONE bounded recovery: the same messages with
|
||
``agent.CORRECTION_INSTRUCTION`` folded into the single system prompt,
|
||
``tools=None``, a fresh filter, the same phase-67 retry budget, streamed
|
||
through the same piece loop (extracted as the inner ``_pump`` helper).
|
||
A recovery that also comes back empty — and any grounded round where the
|
||
recovery policy in ``run_agent`` fails — settles with the dedicated
|
||
structured ``error`` frame (``MalformedReplyError`` caught before the
|
||
generic ``LLMError`` handler): no ``done``, no ``query_log`` row,
|
||
byte-for-byte the existing terminal-error shape. The per-turn log line
|
||
records ``scaffold_stripped=N`` after ``retries=N`` — the sum across the
|
||
turn's requests (grounded: the agent's rounds + forced final + any
|
||
recovery, via the holder; deflected: this turn's filters), 0 on clean
|
||
turns (the field is uniform, the phase-67 ``retries=N`` pattern); the
|
||
recovery does not bump ``retries=N`` (it is not a phase-67
|
||
endpoint-retry).
|
||
|
||
Chat history (phase 74, TODO L4, owner-locked A2/A3/A4 2026-09-08):
|
||
``POST /api/chat`` accepts an optional ``history`` — the client's prior
|
||
turns, oldest first (the ``bor.chat.v1`` record minus the current
|
||
question; the endpoint stays stateless per A10 — nothing is stored). It
|
||
is mapped ONCE per turn by :func:`app.rag.prompts.history_to_messages`
|
||
— trimmed newest-first against the settings budgets
|
||
(``history_max_turns`` / ``history_max_chars``; a capped-out turn is
|
||
dropped whole, never truncated) — and fed to the model on BOTH turn
|
||
branches: the deflected path splices it between the system prompt and
|
||
the current user message (the phase-71 recovery still rebuilds from
|
||
``messages[1:]`` — unchanged), and the grounded agent receives it as
|
||
``run_agent(..., history=hist)``. Prior brain turns' thinking travels
|
||
as ``reasoning_content`` on the assistant message (the preserve-
|
||
thinking wire convention, A4). The per-turn log line records
|
||
``history_msgs=N`` after ``kb_chars=N`` (0 when the request carries no
|
||
history — the two-message request stays byte-identical).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
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.core.auth import require_user
|
||
from app.db import SessionLocal, db_available
|
||
from app.models import Document, QueryLog
|
||
from app.rag.agent import (
|
||
CORRECTION_INSTRUCTION, # phase 71: the harness-owned recovery line
|
||
AgentHolder,
|
||
MalformedReplyError, # phase 71: the recovery policy's terminal signal
|
||
run_agent,
|
||
)
|
||
from app.rag.llm import (
|
||
EmbeddingError,
|
||
LLMClient,
|
||
LLMError,
|
||
RetryPiece, # phase 67: one LLM request restart (an SSE retry frame)
|
||
StreamPiece, # type of the answer pieces streamed by the agent loop
|
||
ToolCallPiece, # phase 37: one model-requested tool call
|
||
ToolResultPiece, # phase 95: one truncated tool result (SSE frame = task 02)
|
||
chat_stream_retried, # phase 67: the retry-before-first-piece primitive
|
||
)
|
||
from app.rag.overview import load_kb_overview
|
||
from app.rag.prompts import build_deflect_prompt, build_high_prompt, history_to_messages
|
||
from app.rag.retriever import RetrievedChunk, retrieve, select_documents, weak_hit_titles
|
||
from app.rag.scaffolding import ScaffoldingFilter # phase 71: the streaming filter
|
||
from app.rag.suggestions import derive_suggestions
|
||
from app.schemas import (
|
||
ChatDoneEvent,
|
||
ChatErrorEvent,
|
||
ChatRequest,
|
||
ChatRetryEvent,
|
||
ChatThinkingEvent,
|
||
ChatToolEvent,
|
||
ChatToolResultEvent,
|
||
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"}
|
||
|
||
#: Concurrency cap for /api/chat (SEC-14-04, task 03). A module-level
|
||
#: counter for the fast-path pre-check and a Semaphore for the real gate.
|
||
_chat_active: int = 0
|
||
_chat_semaphore: asyncio.Semaphore | None = None
|
||
|
||
|
||
def _get_chat_semaphore() -> asyncio.Semaphore:
|
||
"""Lazy-initialize the chat concurrency semaphore (SEC-14-04).
|
||
|
||
The semaphore is initialized on first use so that ``get_settings()``
|
||
is called with the correct environment — never at import time.
|
||
"""
|
||
global _chat_semaphore
|
||
if _chat_semaphore is None:
|
||
settings = get_settings()
|
||
_chat_semaphore = asyncio.Semaphore(max(1, settings.chat_max_concurrent))
|
||
return _chat_semaphore
|
||
|
||
|
||
_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,
|
||
_user: None = Depends(require_user), # noqa: B008 # phase 79: admin or live token
|
||
llm: LLMClient = Depends(get_llm), # noqa: B008
|
||
):
|
||
"""One chat turn: SSE stream of ``delta`` events + a final ``done``.
|
||
|
||
User-gated (phase 79): anonymous callers get 401 ``authentication
|
||
required`` before any streaming — the ONLY anonymous content is the
|
||
shared chats.
|
||
|
||
DB sessions (SEC-14-04): no long-lived session is held across the
|
||
SSE stream — every DB step (retrieval, tool execution, query_log
|
||
write) uses its own short-lived session via ``SessionLocal()``.
|
||
|
||
Concurrency (SEC-14-04, task 03): a semaphore limits concurrent
|
||
turns; excess requests get a 503 error.
|
||
"""
|
||
# SEC-14-04 / task 03: concurrency cap — reject immediately if at
|
||
# capacity (the fast-path pre-check; the semaphore inside stream()
|
||
# is the real gate that prevents races).
|
||
settings = get_settings()
|
||
max_concurrent = settings.chat_max_concurrent
|
||
if _chat_active >= max_concurrent:
|
||
return JSONResponse(
|
||
status_code=503,
|
||
content={"detail": "Too many concurrent chat turns — try again."},
|
||
)
|
||
|
||
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()
|
||
|
||
def db_factory() -> Session:
|
||
"""SEC-14-04: session factory for short-lived sessions."""
|
||
return SessionLocal()
|
||
|
||
async def stream() -> AsyncIterator[str]:
|
||
# SEC-14-04 / task 03: concurrency accounting — increment before
|
||
# the semaphore acquire (the pre-check counter is best-effort;
|
||
# the semaphore is the real gate). The outer finally decrements
|
||
# when the stream ends (normal or error).
|
||
global _chat_active
|
||
_chat_active += 1
|
||
sem = _get_chat_semaphore()
|
||
await sem.acquire()
|
||
try:
|
||
# 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
|
||
retries_used = 0 # phase 67: LLM requests restarted this turn (log line)
|
||
try:
|
||
settings = get_settings()
|
||
# Phase 74 (TODO L4): the client's prior turns, mapped ONCE
|
||
# per turn — trimmed newest-first against the settings
|
||
# budgets, assistant turns carrying their prior thinking as
|
||
# ``reasoning_content`` (A4). BOTH branches below (deflected
|
||
# + grounded agent) reuse the same block; an absent/empty
|
||
# history yields ``[]`` (the byte-identical two-message
|
||
# request, A2).
|
||
hist = history_to_messages(request.history, settings)
|
||
|
||
# 1. Embed the question.
|
||
# Phase 67: a dead embeddings endpoint is retried before any
|
||
# frame has left the server — up to ``llm_retries`` restarts,
|
||
# a flat ``llm_retry_delay`` between attempts, one SSE
|
||
# ``retry`` frame per restart (the UI shows the transient
|
||
# "retrying" status, not an error — locked A4). The final
|
||
# failure keeps the EXISTING terminal ``error`` frame (the
|
||
# copy reads correctly after N tries); ``llm_retries=0`` is
|
||
# byte-identical to the pre-phase-67 single attempt.
|
||
t0 = time.monotonic()
|
||
max_attempts = settings.llm_retries + 1
|
||
attempt = 1
|
||
while True:
|
||
try:
|
||
question_vec = await llm.embed_one(request.message)
|
||
break
|
||
except EmbeddingError as e:
|
||
embed_ms = int((time.monotonic() - t0) * 1000)
|
||
if attempt >= max_attempts:
|
||
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
|
||
logger.warning(
|
||
"chat: question=%r embedding failed (attempt %d/%d) — "
|
||
"retrying in %.1fs: %s",
|
||
request.message,
|
||
attempt,
|
||
max_attempts,
|
||
settings.llm_retry_delay,
|
||
e,
|
||
)
|
||
retries_used += 1
|
||
yield sse_event(
|
||
ChatRetryEvent(
|
||
attempt=attempt + 1, max_attempts=max_attempts
|
||
).model_dump()
|
||
)
|
||
await asyncio.sleep(settings.llm_retry_delay)
|
||
attempt += 1
|
||
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.
|
||
# SEC-14-04: each step uses a short-lived session.
|
||
try:
|
||
with SessionLocal() as step_db:
|
||
steering_notes = load_steering_notes(step_db)
|
||
with SessionLocal() as step_db:
|
||
kb_overview = load_kb_overview(step_db)
|
||
with SessionLocal() as step_db:
|
||
chunks = retrieve(step_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: list[dict[str, Any]] = [
|
||
{"role": "system", "content": plan.system_prompt},
|
||
*hist, # phase 74: the trimmed prior turns (empty by default)
|
||
{"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 | RetryPiece | ToolResultPiece
|
||
]
|
||
deflected_filter: ScaffoldingFilter | None = None
|
||
if plan.deflected:
|
||
# Phase 67: the deflected stream goes through the retry
|
||
# primitive — a dead endpoint is restarted (SSE ``retry``
|
||
# frames) only before its first piece (locked A2); the
|
||
# grounded path stays a plain ``run_agent`` call (task 03
|
||
# makes IT retry internally) — its ``RetryPiece``s flow
|
||
# through the shared piece loop below. Phase 71: the
|
||
# request's content also runs through a caller-owned
|
||
# filter (one per request) — a scaffolding-only reply
|
||
# streams zero ``delta`` frames instead of raw tokens, and
|
||
# the filter's ``stripped_chars`` drives the recovery
|
||
# decision after the piece loop.
|
||
deflected_filter = ScaffoldingFilter()
|
||
answer_stream = chat_stream_retried(
|
||
llm,
|
||
messages,
|
||
tools=None,
|
||
retries=settings.llm_retries,
|
||
delay=settings.llm_retry_delay,
|
||
scaffolding=deflected_filter,
|
||
)
|
||
else:
|
||
answer_stream = run_agent(
|
||
llm,
|
||
db_factory, # SEC-14-04: session factory, not a long-lived session
|
||
system_prompt=plan.system_prompt,
|
||
user_message=request.message,
|
||
seed_docs=plan.docs,
|
||
settings=settings,
|
||
holder=holder,
|
||
history=hist, # phase 74: the same trimmed prior turns
|
||
)
|
||
thinking_chars = 0
|
||
content_chars = 0 # phase 71: the turn's visible (clean) content
|
||
scaffold_stripped = 0 # phase 71: sum across the turn's requests
|
||
|
||
async def _pump(
|
||
pieces: AsyncIterator[
|
||
StreamPiece | ToolCallPiece | RetryPiece | ToolResultPiece
|
||
],
|
||
) -> AsyncIterator[str]:
|
||
"""One request's piece loop (phase 71 extraction): the
|
||
thinking/tool/retry/tool_result/delta handling shared by
|
||
the turn's first pass and — deflected path only — the one
|
||
bounded recovery. Behavior-preserving for the first pass
|
||
(pinned by the existing integration suite). Phase 95:
|
||
the ``ToolResultPiece`` branch emits the additive
|
||
``tool_result`` SSE frame (the seventh, optional event
|
||
type — the A15 extension)."""
|
||
nonlocal thinking_chars, content_chars, retries_used
|
||
async for piece in pieces: # StreamPiece | ToolCallPiece | RetryPiece
|
||
if isinstance(piece, ToolCallPiece):
|
||
# Phase 37 (PLAN §4 extension; phase 70): one SSE
|
||
# ``tool`` frame per model-requested call.
|
||
# ``argument`` is the single string the model
|
||
# passed — ``read``'s ``path`` (the combined
|
||
# ``source/path``), ``grep``'s ``pattern``,
|
||
# ``ls``'s ``path`` — or null (a non-string value
|
||
# is a model error the backend refuses, as is an
|
||
# omitted argument).
|
||
argument = piece.arguments.get(
|
||
"pattern" if piece.name == "grep" else "path"
|
||
)
|
||
argument = argument if isinstance(argument, str) else None
|
||
yield sse_event(
|
||
ChatToolEvent(name=piece.name, argument=argument).model_dump()
|
||
)
|
||
continue
|
||
if isinstance(piece, RetryPiece):
|
||
# Phase 67: the answer stream was restarted before
|
||
# its first piece (locked A2) — a transient status
|
||
# frame, never an error. No other state changes:
|
||
# the thinking/clock/timeout handling is the
|
||
# client's job.
|
||
retries_used += 1
|
||
yield sse_event(
|
||
ChatRetryEvent(
|
||
attempt=piece.attempt, max_attempts=piece.max_attempts
|
||
).model_dump()
|
||
)
|
||
continue
|
||
if isinstance(piece, ToolResultPiece):
|
||
# Phase 95 (A15 extension, task 02): one optional
|
||
# ``tool_result`` frame per truncated ``read`` —
|
||
# emitted HERE, right where the agent loop yielded
|
||
# the piece: AFTER the matching ``tool`` frame and
|
||
# BEFORE the next model round. Additive: a
|
||
# non-truncated read yields no piece at all (no
|
||
# frame), and the other six event types are
|
||
# byte-identical.
|
||
yield sse_event(
|
||
ChatToolResultEvent(
|
||
name=piece.name,
|
||
argument=piece.argument,
|
||
truncated=piece.truncated,
|
||
chars_shown=piece.chars_shown,
|
||
chars_total=piece.chars_total,
|
||
).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:
|
||
content_chars += len(piece.text)
|
||
yield sse_event({"type": "delta", "text": piece.text})
|
||
|
||
try:
|
||
async for frame in _pump(answer_stream):
|
||
yield frame
|
||
if plan.deflected and deflected_filter is not None:
|
||
scaffold_stripped = deflected_filter.stripped_chars
|
||
# Phase 71: the deflected reply's visible content was
|
||
# wiped by the filter (the scaffolding was the whole
|
||
# "answer") — the ONE bounded recovery: the same
|
||
# messages with the correction folded into the single
|
||
# system prompt, ``tools=None``, a FRESH filter, the
|
||
# same phase-67 retry budget, streamed through the
|
||
# same piece loop. A round with real visible content
|
||
# needs no recovery (the clean content stands).
|
||
if content_chars == 0 and deflected_filter.stripped_chars > 0:
|
||
logger.warning(
|
||
"chat: deflected reply was pure tool-scaffolding "
|
||
"(%d chars stripped) — running the one bounded "
|
||
"recovery",
|
||
deflected_filter.stripped_chars,
|
||
)
|
||
recovery_filter = ScaffoldingFilter()
|
||
recovery_stream = chat_stream_retried(
|
||
llm,
|
||
[
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
plan.system_prompt + "\n"
|
||
+ CORRECTION_INSTRUCTION
|
||
),
|
||
},
|
||
*messages[1:],
|
||
],
|
||
tools=None,
|
||
retries=settings.llm_retries,
|
||
delay=settings.llm_retry_delay,
|
||
scaffolding=recovery_filter,
|
||
)
|
||
async for frame in _pump(recovery_stream):
|
||
yield frame
|
||
scaffold_stripped += recovery_filter.stripped_chars
|
||
if content_chars == 0:
|
||
# The second empty reply is terminal (at most
|
||
# one recovery per turn) — the dedicated error
|
||
# frame below (no done, no query_log row).
|
||
logger.warning(
|
||
"chat: the recovery reply was still empty "
|
||
"(scaffold_stripped=%d) — settling with a "
|
||
"malformed-reply error",
|
||
scaffold_stripped,
|
||
)
|
||
raise MalformedReplyError(
|
||
"the deflected model answered in raw "
|
||
"tool-scaffolding twice in a row — no "
|
||
"clean answer to stream"
|
||
)
|
||
else:
|
||
# Grounded turns: the agent's rounds + forced final +
|
||
# any recovery already accumulated the turn total on
|
||
# the holder (the deflected fallback is 0 — the agent
|
||
# never runs, so this branch is grounded-only).
|
||
scaffold_stripped = holder.scaffold_stripped
|
||
except MalformedReplyError as e:
|
||
# Phase 71: the recovery policy's terminal signal —
|
||
# caught BEFORE the generic LLMError handler (it
|
||
# subclasses it), so the dedicated copy reaches the UI;
|
||
# the generic "dropped the connection" copy stays for
|
||
# transport failures.
|
||
logger.error(
|
||
"chat: malformed reply after the one bounded recovery "
|
||
"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 model returned a malformed reply — please try again."
|
||
).model_dump()
|
||
)
|
||
return
|
||
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 (the drill-down ``ls`` /
|
||
# ``read`` / ``grep`` lookups) 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:
|
||
with SessionLocal() as log_db:
|
||
log_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,
|
||
)
|
||
)
|
||
log_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 history_msgs=%d threshold=%.2f deflected=%s sources=%r "
|
||
"thinking_chars=%d tool_calls=%d total_ms=%d retries=%d scaffold_stripped=%d",
|
||
request.message,
|
||
embed_ms,
|
||
plan.top_score,
|
||
plan.fts_hits,
|
||
plan.summary_hits,
|
||
plan.tuning_count,
|
||
plan.kb_chars,
|
||
len(hist),
|
||
settings.relevance_threshold,
|
||
plan.deflected,
|
||
source_paths,
|
||
thinking_chars,
|
||
holder.tool_calls,
|
||
total_ms,
|
||
retries_used,
|
||
scaffold_stripped,
|
||
)
|
||
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),
|
||
)
|
||
finally:
|
||
sem.release()
|
||
_chat_active -= 1
|
||
|
||
return StreamingResponse(stream(), media_type="text/event-stream", headers=SSE_HEADERS)
|