"""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", }, )