feat(chat): stop an in-flight answer — Send becomes Stop, the partial is kept and persisted, the model stream is torn down
This commit is contained in:
+207
-176
@@ -240,191 +240,222 @@ async def chat(
|
||||
started = time.monotonic()
|
||||
|
||||
async def stream() -> AsyncIterator[str]:
|
||||
# 1. Embed the question.
|
||||
t0 = time.monotonic()
|
||||
# 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:
|
||||
question_vec = await llm.embed_one(request.message)
|
||||
except EmbeddingError as e:
|
||||
# 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)
|
||||
logger.error(
|
||||
"chat: question=%r embed_ms=%d total_ms=%d — embedding failed: %s",
|
||||
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,
|
||||
e,
|
||||
)
|
||||
settled = True # terminal: the done 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),
|
||||
)
|
||||
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,
|
||||
)
|
||||
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),
|
||||
)
|
||||
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).
|
||||
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),
|
||||
ChatDoneEvent(
|
||||
deflected=plan.deflected,
|
||||
sources=", ".join(source_paths),
|
||||
latency_ms=total_ms,
|
||||
)
|
||||
sources=[
|
||||
SourceRef(source=d.source, path=d.path, title=d.title) for d in cited_docs
|
||||
],
|
||||
suggestions=plan.suggestions,
|
||||
).model_dump()
|
||||
)
|
||||
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,
|
||||
)
|
||||
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)
|
||||
|
||||
+22
-10
@@ -246,12 +246,19 @@ async def run_agent(
|
||||
rounds = 0
|
||||
while True:
|
||||
calls: list[ToolCallPiece] = []
|
||||
async for piece in llm.chat_stream(
|
||||
cast("list[dict[str, str]]", messages), tools=tools
|
||||
):
|
||||
if isinstance(piece, ToolCallPiece):
|
||||
calls.append(piece)
|
||||
yield piece
|
||||
# Phase 48: bind the round's stream so a consumer abandon
|
||||
# (GeneratorExit into the yield below) tears down the in-flight
|
||||
# model stream deterministically — not GC-dependent. Awaiting
|
||||
# ``aclose()`` in the ``finally`` is safe because it does not
|
||||
# yield; on a fully consumed round it is a quiet no-op.
|
||||
stream = llm.chat_stream(cast("list[dict[str, str]]", messages), tools=tools)
|
||||
try:
|
||||
async for piece in stream:
|
||||
if isinstance(piece, ToolCallPiece):
|
||||
calls.append(piece)
|
||||
yield piece
|
||||
finally:
|
||||
await stream.aclose()
|
||||
if not calls:
|
||||
return # the answer was streamed
|
||||
call = calls[0] # a stream can carry several calls; run the first
|
||||
@@ -287,8 +294,13 @@ async def run_agent(
|
||||
"no-tools answer",
|
||||
rounds,
|
||||
)
|
||||
async for piece in llm.chat_stream(
|
||||
cast("list[dict[str, str]]", messages), tools=None
|
||||
):
|
||||
yield piece
|
||||
# Phase 48: the forced final answer gets the same explicit
|
||||
# teardown as the loop rounds (consumer abandon mid-final
|
||||
# answer must still close the model's stream).
|
||||
final = llm.chat_stream(cast("list[dict[str, str]]", messages), tools=None)
|
||||
try:
|
||||
async for piece in final:
|
||||
yield piece
|
||||
finally:
|
||||
await final.aclose()
|
||||
return
|
||||
|
||||
+28
-5
@@ -20,12 +20,12 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncGenerator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from openai import AsyncOpenAI, AsyncStream
|
||||
from openai.types.chat import ChatCompletionChunk, ChatCompletionMessageParam
|
||||
|
||||
from app.config import Settings, get_settings
|
||||
|
||||
@@ -318,7 +318,7 @@ class LLMClient:
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> AsyncIterator[StreamPiece | ToolCallPiece]:
|
||||
) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
|
||||
"""Stream assistant pieces from the chat model (PLAN A5/A15, phase 17).
|
||||
|
||||
``stream=True`` against the OpenAI-compatible endpoint, yielding
|
||||
@@ -355,6 +355,17 @@ class LLMClient:
|
||||
Any failure (network, HTTP, malformed stream) surfaces as
|
||||
:class:`LLMError` so the API layer can turn it into an SSE
|
||||
``error`` event instead of a hung request.
|
||||
|
||||
Teardown (phase 48, 2026-08-29, ``TODO.md`` L3): once
|
||||
``create()`` succeeded, the endpoint stream's lifetime is
|
||||
explicit — it is closed on **every** exit: normal exhaustion
|
||||
(a quiet no-op on the already-ended SDK stream, so the
|
||||
completed path stays byte-identical), a wrapped mid-stream
|
||||
failure, and consumer abandon (stop/cancel — ``GeneratorExit``;
|
||||
awaiting in the ``finally`` is safe because it does not yield).
|
||||
The SDK's ``close()`` awaits the underlying httpx response's
|
||||
``aclose()``, so the local model stops generating as soon as
|
||||
the SSE consumer goes away.
|
||||
"""
|
||||
kwargs: dict[str, Any] = {
|
||||
# ``{role, content}`` dicts are exactly what the message params
|
||||
@@ -367,8 +378,12 @@ class LLMClient:
|
||||
}
|
||||
if tools is not None:
|
||||
kwargs["tools"] = tools
|
||||
stream: AsyncStream[ChatCompletionChunk] | None = None
|
||||
try:
|
||||
stream = await self._client.chat.completions.create(**kwargs)
|
||||
stream = cast(
|
||||
"AsyncStream[ChatCompletionChunk]",
|
||||
await self._client.chat.completions.create(**kwargs),
|
||||
)
|
||||
calls: dict[int, _ToolCallSlot] = {}
|
||||
emitted = False
|
||||
async for chunk in stream:
|
||||
@@ -417,6 +432,14 @@ class LLMClient:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 — wrap transport-level failures
|
||||
raise LLMError(f"chat stream from {self.settings.llm_base_url} failed: {e}") from e
|
||||
finally:
|
||||
# Phase 48: deterministic teardown — whenever ``create()``
|
||||
# succeeded, close the endpoint's stream on every subsequent
|
||||
# exit (normal exhaustion, wrapped failures, and consumer
|
||||
# abandon). A failure of ``create()`` itself never sets
|
||||
# ``stream``, so it stays the plain wrap above.
|
||||
if stream is not None:
|
||||
await stream.close()
|
||||
|
||||
|
||||
async def check_models(llm: LLMClient) -> None:
|
||||
|
||||
Reference in New Issue
Block a user