gongzhiyongandClaude Sonnet 4.6
47921d1a45
## Fixes - [#3] Status events now add/update a status ActivityNode in timeline (upsertWsStatusNode) - [#4] Done callback adds done node; error callback adds error node + sets session.status="error" - [#5] handleRegenerate fully synced with workspace: status/tool/card/done/error all handled - [#1][#2] localStorage persistence: completed/error sessions auto-saved, lazy-loaded on demand - [#1] handleSelectConversation preloads workspace sessions from localStorage for history messages - Refactored completeWsSession to include done ActivityNode in timeline - Added errorWsSession, upsertWsStatusNode, saveWsToStorage, loadWsFromStorage helpers ## Known limitation - [#7] workspace_card merge:true not yet used by backend (all cards are append-only for now) - History workspace recovery depends on localStorage (browser-local, not cross-device) Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
358 lines
14 KiB
Python
358 lines
14 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.mcp_jina import create_jina_mcp_client, get_jina_mcp_tools
|
||
|
||
|
||
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 _run_graph_stream(
|
||
request: ChatRequest,
|
||
graph,
|
||
all_tools: list,
|
||
) -> AsyncIterator[bytes]:
|
||
"""Run the LangGraph agent and yield SSE events."""
|
||
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 all_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
|
||
|
||
try:
|
||
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):
|
||
full_content.append(chunk.content)
|
||
yield _sse({"type": "token", "content": chunk.content})
|
||
|
||
elif kind == "on_tool_start":
|
||
call_id = event.get("run_id", str(uuid.uuid4()))
|
||
tool_name = event.get("name", "unknown")
|
||
tool_input = event.get("data", {}).get("input", {})
|
||
ts = int(time.time() * 1000)
|
||
tool_start_ts[call_id] = ts
|
||
has_tool_activity = True
|
||
yield _sse({
|
||
"type": "tool_start",
|
||
"call_id": call_id,
|
||
"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":
|
||
call_id = event.get("run_id", "")
|
||
tool_name = event.get("name", "unknown")
|
||
output = event.get("data", {}).get("output", "")
|
||
output_str = _extract_output_str(output)
|
||
ts = int(time.time() * 1000)
|
||
duration_ms = ts - tool_start_ts.pop(call_id, ts)
|
||
if _is_tool_error(output_str):
|
||
yield _sse({
|
||
"type": "tool_error",
|
||
"call_id": call_id,
|
||
"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",
|
||
"call_id": call_id,
|
||
"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,
|
||
})
|
||
|
||
elif kind == "on_chain_start":
|
||
chain_name = event.get("name", "")
|
||
tags = event.get("tags", [])
|
||
is_graph_step = any(t.startswith("graph:step:") for t in tags)
|
||
if chain_name == "agent" and is_graph_step:
|
||
if has_tool_activity:
|
||
yield _sse({
|
||
"type": "status",
|
||
"stage": "generating",
|
||
"message": "正在生成回复...",
|
||
})
|
||
else:
|
||
yield _sse({
|
||
"type": "status",
|
||
"stage": "thinking",
|
||
"message": "正在分析...",
|
||
})
|
||
|
||
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"})
|
||
|
||
|
||
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
|
||
regular_tools, needs_jina_mcp = resolve_tools(request.tools)
|
||
|
||
if needs_jina_mcp:
|
||
mcp_client = create_jina_mcp_client()
|
||
jina_tools = await get_jina_mcp_tools(mcp_client)
|
||
all_tools = regular_tools + jina_tools
|
||
graph = await get_chat_graph(model=request.model, tools=all_tools)
|
||
async for chunk in _run_graph_stream(request, graph, all_tools):
|
||
yield chunk
|
||
else:
|
||
graph = await get_chat_graph(model=request.model, tools=regular_tools)
|
||
async for chunk in _run_graph_stream(request, graph, regular_tools):
|
||
yield chunk
|
||
|
||
|
||
# ── Trace helpers ─────────────────────────────────────────────────────────────
|
||
|
||
def _extract_output_str(output) -> str:
|
||
"""Extract plain text from tool output, handling MCP structured content."""
|
||
if isinstance(output, str):
|
||
return output
|
||
# MCP ToolMessage: has .content as list of {'type': 'text', 'text': '...'}
|
||
content = getattr(output, "content", None)
|
||
if isinstance(content, list):
|
||
parts = [c.get("text", "") for c in content if isinstance(c, dict) and c.get("type") == "text"]
|
||
return "\n".join(parts)
|
||
if isinstance(content, str):
|
||
return content
|
||
return str(output)
|
||
|
||
|
||
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": "查询工单详情",
|
||
"generate_document": "生成文档",
|
||
"sandbox_run": "执行沙盒代码",
|
||
# Google Serper
|
||
"serper_search": "Google 搜索",
|
||
# Jina MCP tools
|
||
"search_web": "外部搜索",
|
||
"read_url": "读取网页",
|
||
"sort_by_relevance": "相关性排序",
|
||
}
|
||
|
||
_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 "serper_search":
|
||
return f"搜索：{str(inp.get('query', ''))[:50]}"
|
||
case "search_web":
|
||
return f"搜索：{str(inp.get('query', ''))[:50]}"
|
||
case "read_url":
|
||
url = inp.get('url', '')
|
||
if isinstance(url, list):
|
||
url = url[0] if url else ''
|
||
return f"读取：{str(url)[:60]}"
|
||
case "sort_by_relevance":
|
||
query = str(inp.get('query', ''))[:30]
|
||
count = len(inp.get('documents', []))
|
||
return f"排序 {count} 条结果，查询：{query}"
|
||
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 "serper_search":
|
||
return f"Google 搜索完成，{len(output.splitlines())} 行结果"
|
||
case "search_web":
|
||
count = output.count("##")
|
||
return f"找到 {max(count, 1)} 条搜索结果"
|
||
case "read_url":
|
||
lines = len(output.splitlines())
|
||
return f"读取完成，{lines} 行内容"
|
||
case "sort_by_relevance":
|
||
return "排序完成"
|
||
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",
|
||
},
|
||
)
|