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)
|
||||
|
||||
Reference in New Issue
Block a user