phase: 110_fix_sse_db_pool_exhaustion
--- **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/`.
This commit is contained in:
+467
-405
@@ -166,7 +166,7 @@ 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 db_available, get_db
|
||||
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
|
||||
@@ -206,6 +206,25 @@ 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
|
||||
|
||||
|
||||
@@ -316,7 +335,6 @@ def plan_turn(
|
||||
async def chat(
|
||||
request: ChatRequest,
|
||||
_user: None = Depends(require_user), # noqa: B008 # phase 79: admin or live token
|
||||
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``.
|
||||
@@ -324,7 +342,25 @@ async def chat(
|
||||
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,
|
||||
@@ -338,423 +374,449 @@ async def chat(
|
||||
|
||||
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
|
||||
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)
|
||||
def db_factory() -> Session:
|
||||
"""SEC-14-04: session factory for short-lived sessions."""
|
||||
return SessionLocal()
|
||||
|
||||
# 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",
|
||||
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,
|
||||
embed_ms,
|
||||
total_ms,
|
||||
attempt,
|
||||
max_attempts,
|
||||
settings.llm_retry_delay,
|
||||
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.
|
||||
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: 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,
|
||||
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
|
||||
attempt=attempt + 1, max_attempts=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})
|
||||
await asyncio.sleep(settings.llm_retry_delay)
|
||||
attempt += 1
|
||||
embed_ms = int((time.monotonic() - t0) * 1000)
|
||||
|
||||
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:
|
||||
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,
|
||||
# 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
|
||||
)
|
||||
)
|
||||
db.commit()
|
||||
except Exception: # noqa: BLE001 — the answer already went out
|
||||
logger.exception("chat: failed to write query_log question=%r", request.message)
|
||||
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},
|
||||
]
|
||||
|
||||
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
|
||||
# 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
|
||||
],
|
||||
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",
|
||||
) -> 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,
|
||||
int((time.monotonic() - started) * 1000),
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user