gongzhiyongandClaude Sonnet 4.6
722ec3e81d
Add CoT execution trace to the SSE stream: status events for stage transitions, enriched tool_start with input_summary, tool_end with output_summary and duration_ms, and tool_error for failed tool calls. Helper functions _sse, _summarize_input, _summarize_output provide human-readable summaries for each tool type. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
296 lines
11 KiB
Python
296 lines
11 KiB
Python
"""SSE streaming chat endpoint."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import time
|
||
import uuid
|
||
from collections.abc import AsyncIterator
|
||
|
||
from langchain_core.messages import HumanMessage
|
||
from litestar import post
|
||
from litestar.response import Stream
|
||
|
||
from app.graph.builder import get_chat_graph
|
||
from app.schemas import ChatRequest
|
||
from app.store.memory import get_checkpointer
|
||
from app.store.postgres import Conversation, Message, async_session_factory
|
||
from app.tools import resolve_tools
|
||
from app.tools.search import set_search_model
|
||
|
||
|
||
async def _ensure_conversation(conversation_id: str, first_message: str) -> None:
|
||
"""Create conversation and persist the user message."""
|
||
async with async_session_factory() as session:
|
||
existing = await session.get(Conversation, conversation_id)
|
||
if existing is None:
|
||
# Use first ~50 chars of message as title
|
||
title = first_message[:50].strip() or "New conversation"
|
||
conv = Conversation(id=conversation_id, title=title)
|
||
session.add(conv)
|
||
# Persist user message
|
||
msg = Message(
|
||
conversation_id=conversation_id,
|
||
role="human",
|
||
content=first_message,
|
||
)
|
||
session.add(msg)
|
||
await session.commit()
|
||
|
||
|
||
async def _persist_ai_message(conversation_id: str, content: str) -> None:
|
||
"""Persist the AI response message."""
|
||
async with async_session_factory() as session:
|
||
msg = Message(
|
||
conversation_id=conversation_id,
|
||
role="ai",
|
||
content=content,
|
||
)
|
||
session.add(msg)
|
||
await session.commit()
|
||
|
||
|
||
async def _stream_response(request: ChatRequest) -> AsyncIterator[bytes]:
|
||
"""Stream LLM response tokens via SSE."""
|
||
# Ensure conversation exists and persist user message
|
||
await _ensure_conversation(request.conversation_id, request.message)
|
||
|
||
# Resolve tools from frontend tool keys
|
||
active_tools = resolve_tools(request.tools)
|
||
|
||
# Set search strategy model so the web_search tool knows depth
|
||
set_search_model(request.model)
|
||
|
||
graph = await get_chat_graph(model=request.model, tools=active_tools)
|
||
|
||
config = {
|
||
"configurable": {"thread_id": request.conversation_id},
|
||
}
|
||
|
||
# When using ReAct agent (with tools), input is just messages.
|
||
# When using plain graph (no tools), input includes model key.
|
||
if active_tools:
|
||
input_data = {"messages": [HumanMessage(content=request.message)]}
|
||
else:
|
||
input_data = {
|
||
"messages": [HumanMessage(content=request.message)],
|
||
"model": request.model,
|
||
}
|
||
|
||
full_content: list[str] = []
|
||
logger = logging.getLogger(__name__)
|
||
tool_start_ts: dict[str, int] = {} # track per-tool start timestamps
|
||
has_tool_activity = False # whether any tool has been called
|
||
final_status_emitted = False # whether "整理答案" status was emitted
|
||
|
||
try:
|
||
# Emit initial status before graph starts
|
||
yield _sse({"type": "status", "stage": "分析问题", "message": "正在理解您的问题..."})
|
||
|
||
async for event in graph.astream_events(
|
||
input_data,
|
||
config=config,
|
||
version="v2",
|
||
):
|
||
kind = event.get("event", "")
|
||
|
||
if kind == "on_chat_model_stream":
|
||
chunk = event.get("data", {}).get("chunk")
|
||
if chunk and hasattr(chunk, "content") and chunk.content:
|
||
# Only stream text content, skip tool call chunks
|
||
if isinstance(chunk.content, str):
|
||
# Before first token after tool activity, emit "整理答案" status
|
||
if has_tool_activity and not final_status_emitted:
|
||
yield _sse({
|
||
"type": "status",
|
||
"stage": "整理答案",
|
||
"message": "正在结合检索结果生成回复...",
|
||
})
|
||
final_status_emitted = True
|
||
full_content.append(chunk.content)
|
||
yield _sse({"type": "token", "content": chunk.content})
|
||
|
||
elif kind == "on_tool_start":
|
||
tool_name = event.get("name", "unknown")
|
||
tool_input = event.get("data", {}).get("input", {})
|
||
ts = int(time.time() * 1000)
|
||
tool_start_ts[tool_name] = ts
|
||
has_tool_activity = True
|
||
yield _sse({
|
||
"type": "tool_start",
|
||
"tool": tool_name,
|
||
"title": _TOOL_TITLES.get(tool_name, tool_name),
|
||
"input_summary": _summarize_input(tool_name, tool_input),
|
||
"ts": ts,
|
||
})
|
||
|
||
elif kind == "on_tool_end":
|
||
tool_name = event.get("name", "unknown")
|
||
output = event.get("data", {}).get("output", "")
|
||
output_str = output if isinstance(output, str) else str(output)
|
||
ts = int(time.time() * 1000)
|
||
duration_ms = ts - tool_start_ts.pop(tool_name, ts)
|
||
if _is_tool_error(output_str):
|
||
yield _sse({
|
||
"type": "tool_error",
|
||
"tool": tool_name,
|
||
"title": _TOOL_TITLES.get(tool_name, tool_name),
|
||
"error_summary": _extract_error_summary(output_str),
|
||
"duration_ms": duration_ms,
|
||
"ts": ts,
|
||
})
|
||
else:
|
||
yield _sse({
|
||
"type": "tool_end",
|
||
"tool": tool_name,
|
||
"title": _TOOL_TITLES.get(tool_name, tool_name),
|
||
"output_summary": _summarize_output(tool_name, output_str),
|
||
"status": "success",
|
||
"duration_ms": duration_ms,
|
||
"ts": ts,
|
||
})
|
||
|
||
except Exception as exc:
|
||
logger.error("SSE stream error for conversation %s: %s", request.conversation_id, exc, exc_info=True)
|
||
|
||
error_msg = str(exc)
|
||
user_hint = "请求处理出现错误，请重试"
|
||
|
||
# Detect checkpoint pollution: a prior tool crash left an AIMessage
|
||
# with tool_calls but no corresponding ToolMessage. LangGraph refuses
|
||
# to continue the thread. Purge the thread so the next request
|
||
# starts from a clean state.
|
||
if "tool_calls" in error_msg and "ToolMessage" in error_msg:
|
||
try:
|
||
checkpointer = await get_checkpointer()
|
||
await checkpointer.adelete_thread(request.conversation_id)
|
||
logger.warning(
|
||
"Purged polluted checkpoint for thread %s",
|
||
request.conversation_id,
|
||
)
|
||
user_hint = "对话状态异常，已自动重置，请重新发送消息"
|
||
except Exception:
|
||
logger.warning(
|
||
"Failed to purge checkpoint for thread %s",
|
||
request.conversation_id,
|
||
exc_info=True,
|
||
)
|
||
|
||
error_data = json.dumps(
|
||
{"type": "error", "content": user_hint},
|
||
ensure_ascii=False,
|
||
)
|
||
yield f"data: {error_data}\n\n".encode("utf-8")
|
||
|
||
finally:
|
||
# Persist whatever AI content was streamed before the error (if any)
|
||
ai_content = "".join(full_content)
|
||
if ai_content:
|
||
await _persist_ai_message(request.conversation_id, ai_content)
|
||
|
||
# Always send done so the frontend closes the stream cleanly
|
||
yield _sse({"type": "done"})
|
||
|
||
|
||
# ── Trace helpers ─────────────────────────────────────────────────────────────
|
||
|
||
def _sse(data: dict) -> bytes:
|
||
return f"data: {json.dumps(data, ensure_ascii=False)}\n\n".encode("utf-8")
|
||
|
||
|
||
_TOOL_TITLES: dict[str, str] = {
|
||
"kb_search": "检索知识库",
|
||
"ticket_list": "查询工单列表",
|
||
"ticket_detail": "查询工单详情",
|
||
"web_search": "外部搜索",
|
||
"generate_document": "生成文档",
|
||
"sandbox_run": "执行沙盒代码",
|
||
}
|
||
|
||
_ERROR_KEYWORDS = (
|
||
"出错", "失败", "超时", "error", "failed", "timeout",
|
||
"not available", "no download link", "execution failed",
|
||
)
|
||
|
||
|
||
def _is_tool_error(output: str) -> bool:
|
||
lo = output.lower()
|
||
return any(kw in lo for kw in _ERROR_KEYWORDS)
|
||
|
||
|
||
def _extract_error_summary(output: str) -> str:
|
||
first_line = output.split("\n")[0].strip()
|
||
return first_line[:80] if first_line else "工具调用失败"
|
||
|
||
|
||
def _summarize_input(tool_name: str, inp: dict | str) -> str:
|
||
if isinstance(inp, str):
|
||
return inp[:60]
|
||
match tool_name:
|
||
case "kb_search":
|
||
return f"查询：{str(inp.get('query', ''))[:50]}"
|
||
case "ticket_list":
|
||
return f"第 {inp.get('page', 1)} 页，每页 {inp.get('page_size', 20)} 条"
|
||
case "ticket_detail":
|
||
return f"工单 ID：{inp.get('ticket_id', '')}"
|
||
case "web_search":
|
||
return f"搜索：{str(inp.get('query', ''))[:50]}"
|
||
case "generate_document":
|
||
return str(inp.get('prompt', ''))[:60]
|
||
case "sandbox_run":
|
||
lang = inp.get("language", "python")
|
||
lines = len(str(inp.get("code", "")).splitlines())
|
||
return f"{lang} 代码（{lines} 行）"
|
||
case _:
|
||
return str(inp)[:60]
|
||
|
||
|
||
def _summarize_output(tool_name: str, output: str) -> str:
|
||
if not output or output.strip() in ("", "(no output)"):
|
||
return "无结果"
|
||
match tool_name:
|
||
case "kb_search":
|
||
count = output.count("---") + 1 if "---" in output else 1
|
||
return f"命中 {count} 条知识库记录"
|
||
case "ticket_list":
|
||
import re as _re
|
||
m = _re.search(r"Found (\d+) tickets", output)
|
||
return f"返回 {m.group(1)} 条工单" if m else "工单列表已获取"
|
||
case "ticket_detail":
|
||
return "工单详情已获取"
|
||
case "web_search":
|
||
count = output.count("##")
|
||
return f"找到 {max(count, 1)} 条搜索结果"
|
||
case "generate_document":
|
||
if "Download:" in output:
|
||
if "[PPT]" in output:
|
||
return "PPT 已生成，可下载"
|
||
if "[Excel]" in output or "[Table]" in output:
|
||
return "表格已生成，可下载"
|
||
return "Word 文档已生成，可下载"
|
||
return "文档生成完成"
|
||
case "sandbox_run":
|
||
lines = len(output.splitlines())
|
||
has_error = output.startswith("[Exit code:")
|
||
return f"执行完成，输出 {lines} 行{'（含错误）' if has_error else ''}"
|
||
case _:
|
||
return output[:60]
|
||
|
||
|
||
@post("/api/chat/stream")
|
||
async def stream_chat(data: ChatRequest) -> Stream:
|
||
"""POST /api/chat/stream - SSE streaming chat endpoint."""
|
||
if not data.conversation_id:
|
||
data.conversation_id = str(uuid.uuid4())
|
||
|
||
return Stream(
|
||
_stream_response(data),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"Connection": "keep-alive",
|
||
"X-Accel-Buffering": "no",
|
||
},
|
||
)
|