"""SSE 流式内核。 职责:把 LangGraph 的 ``astream`` 事件流翻译成 1.5 兼容的 SSE 帧,并在 结束时完成落库、计费、来源标注与上下文计量。 ## 通道分工(重要,避免重复事件) ``stream_mode=["updates", "messages"]`` 两条通道都在推同一个 agent 的进展, 不加区分会导致同一个工具调用被发两次。这里的分工是: - **updates 通道**:工具调用/结果、待办清单、子智能体生命周期、状态提示。 走这里的原因是它给的是**完整的消息对象**(``tool_calls`` 里 args 已解析 完毕、``ToolMessage`` 已成对),不需要在流式增量里拼半个 JSON。 - **messages 通道**:正文 token、思考 token、真实 token 用量。 ``version="v2"`` 下每个 chunk 是 ``{"type", "ns", "data"}``;``subgraphs=True`` 让子智能体的产出带 ``ns``。两个格式都做了兼容(旧版是 ``(mode, data)`` 元组)。 ## 与 1.5 的差异 1.5 的 sse_core 是 1908 行,里面混着计费、压缩、Word 导出、来源刮取等多条 业务线。2.0 把这些拆成独立组件(context_meter / source_tracker / subagent_tracker),内核只负责「读流 → 派事件 → 收尾」。 """ from __future__ import annotations import asyncio import time from dataclasses import dataclass, field from pathlib import Path from typing import Any, AsyncIterator import config from api import concurrency from api.context_meter import ContextMeter from api.sse_events import EventWriter from api.source_tracker import SourceTracker from api.subagent_tracker import SubagentTracker from billing.meter import UsageMeter from db import chat_store from obs import get_logger, get_trace_id logger = get_logger(__name__) MAIN_AGENT_NAME = "deep-agent" MAIN_AGENT_CN = "深度智能体" # 摘要中间件的模型调用也在 messages 通道里推 token,靠这个标记过滤掉, # 否则摘要正文会混进给用户看的回答里 _SUMMARIZATION_SOURCE = "summarization" @dataclass class StreamRequest: """一次流式请求的全部入参。""" agent: Any session_id: str thread_id: str config: dict[str, Any] message: str user_name: str = "default" kind: str = "chat" model: str = "" original_user_message: str = "" trace_id: str = "" slot: concurrency.TaskSlot | None = None started_at: float = field(default_factory=time.time) agent_name: str = MAIN_AGENT_NAME agent_cn_name: str = MAIN_AGENT_CN # 非 None 表示这是一次中断恢复:不组装消息,直接把 Command(resume=...) 喂回图 resume_payload: Any = None def normalize_chunk(chunk: Any) -> tuple[str | None, tuple[Any, ...], Any]: """把不同版本的 astream chunk 归一成 ``(mode, ns, data)``。 兼容三种形态:v2 的 ``{"type","ns","data"}`` 字典、旧版 ``(mode, data)`` 元组、以及无法识别的原样返回(mode=None,上层据此忽略)。 """ if isinstance(chunk, dict) and "data" in chunk and "type" in chunk: return str(chunk["type"]), tuple(chunk.get("ns") or ()), chunk["data"] if isinstance(chunk, tuple) and len(chunk) == 2 and isinstance(chunk[0], str): return chunk[0], (), chunk[1] return None, (), chunk def _message_text(message: Any) -> str: """把消息 content 拍平成纯文本(content 可能是字符串或内容块列表)。""" content = getattr(message, "content", None) if isinstance(content, str): return content if isinstance(content, list): parts: list[str] = [] for block in content: if isinstance(block, str): parts.append(block) elif isinstance(block, dict): text = block.get("text") if isinstance(text, str): parts.append(text) return "".join(parts) return "" def _reasoning_text(message: Any) -> str: """抽取思考链文本(deepseek 的 reasoning_content)。""" extra = getattr(message, "additional_kwargs", None) if isinstance(extra, dict): value = extra.get("reasoning_content") if isinstance(value, str): return value return "" class ChatStream: """一次 /chat 请求的流式执行。""" def __init__(self, request: StreamRequest) -> None: self.req = request self.writer = EventWriter(request.session_id, request.trace_id or get_trace_id()) self.meter = _meter_for_request() self.tracker = SubagentTracker() self.sources = SourceTracker() self.billing = UsageMeter( user_name=request.user_name, session_id=request.session_id, kind=request.kind, model=request.model or config.MODEL, trace_id=request.trace_id, ) # 累计状态 self.answer_parts: list[str] = [] self.tool_calls: list[dict[str, Any]] = [] self.todos: list[dict[str, Any]] = [] self.interrupt_payloads: list[Any] = [] self.finished = False # ======================================================== # 主入口 # ======================================================== async def generate(self) -> AsyncIterator[str]: """产出整个 SSE 流。保证:计费幂等入账、并发位释放、图游标关闭。 这里显式绑定 trace_id 而不是依赖 HTTP 中间件设的 ContextVar: 响应体的迭代在 Starlette 的中间件调用栈之外执行,中间件里绑的值 在生成器里不保证可见。绑定发生在生成器内部,生命周期与本次流一致。 """ from obs import bind_trace_id, unbind_trace_id token = bind_trace_id(self.req.trace_id or None) stream = None try: yield self.writer.agent_start(self.req.agent_name, self.req.agent_cn_name) yield self.writer.thinking("正在理解您的问题...", self.req.agent_cn_name) graph_input = await self._prepare_input() stream = self.req.agent.astream( graph_input, stream_mode=["updates", "messages"], config=self.req.config, subgraphs=True, version="v2", ) async for frame in self._drive(stream): yield frame # 中断检测必须在流结束后做:interrupt 会让流自然结束 async for frame in self._emit_interrupts(): yield frame async for frame in self._completion(): yield frame except asyncio.CancelledError: # 客户端断连:不当作错误,但要把已产出的部分落库并结算 logger.info("客户端断开连接 session=%s", self.req.session_id) async for frame in self._completion(aborted=True): yield frame raise except GeneratorExit: logger.info("生成器被关闭 session=%s", self.req.session_id) async for frame in self._completion(aborted=True): yield frame raise except Exception as exc: async for frame in self._error_frames(exc): yield frame finally: if stream is not None: close = getattr(stream, "aclose", None) if close is not None: try: await close() except Exception: # pragma: no cover - 关闭失败不影响结算 pass self._finalize() unbind_trace_id(token) # ======================================================== # 输入准备 # ======================================================== async def _prepare_input(self) -> Any: """组装喂给图的输入。 两种情形: - **恢复**(``resume_payload`` 非 None):喂 ``Command(resume=...)``, 不组装消息——图会从挂起的那个节点继续。 - **新请求**:热线程(checkpoint 里已有状态)只喂当轮新消息;冷线程 需要从会话历史补种。 补种这一步不能省,也不能对所有线程都做:热线程若把历史重复追加, 上下文会随轮次按 O(N²) 膨胀——这是长会话最容易被忽略的膨胀来源。 """ if self.req.resume_payload is not None: from langgraph.types import Command logger.info("恢复中断 session=%s", self.req.session_id) return Command(resume=self.req.resume_payload) from langchain_core.messages import AIMessage, HumanMessage from agents import agent_memory messages: list[Any] = [] warm = await agent_memory.has_checkpoint(self.req.thread_id) if not warm: history = chat_store.get_messages( self.req.session_id, limit=config.CTX_RECENT_LIMIT ) for row in history: content = row.get("content") or "" if not content: continue if row.get("msg_role") == "user": messages.append(HumanMessage(content=content)) else: messages.append(AIMessage(content=content)) if messages: logger.info( "冷线程补种历史 session=%s 条数=%d", self.req.session_id, len(messages) ) else: logger.info("热线程,仅喂当轮消息 session=%s", self.req.session_id) # 用户消息先落库:即使这一轮流到一半挂了,历史里也得有用户问了什么 chat_store.save_message(self.req.session_id, "user", self.req.message) self.meter.note_text(self.req.message) messages.append(HumanMessage(content=self.req.message)) # 必须是 {"messages": [...]} 字典:直接传 list 会被当成 __start__ 节点的 # 返回值,图会报 InvalidUpdateError: Expected dict return {"messages": messages} # ======================================================== # 流驱动 # ======================================================== async def _drive(self, stream: Any) -> AsyncIterator[str]: """消费 astream,派发事件帧。带墙钟上限。""" deadline = self.req.started_at + config.TASK_WALL_TIMEOUT timed_out = False iterator = stream.__aiter__() while True: if time.time() > deadline: timed_out = True break try: chunk = await asyncio.wait_for( iterator.__anext__(), timeout=max(deadline - time.time(), 1.0) ) except StopAsyncIteration: break except asyncio.TimeoutError: timed_out = True break mode, ns, data = normalize_chunk(chunk) if mode == "updates": for frame in self._handle_updates(data, ns): yield frame elif mode == "messages": for frame in self._handle_messages(data, ns): yield frame # mode 为 None:无法识别的 chunk,忽略而不是猜 if timed_out: raise concurrency.TaskWallTimeoutError( f"任务超过墙钟上限 {config.TASK_WALL_TIMEOUT}s 已被中止" ) # 子智能体可能因为异常/中断没收到结果,收尾时统一关掉,避免前端卡片长转 for run in self.tracker.finish_all(): yield self.writer.progress( f"子智能体 {run.subagent_type} 未正常返回,已中止", subagent=run.subagent_type, subagent_id=run.call_id, ) # -------- updates 通道 -------- def _handle_updates(self, data: Any, ns: tuple[Any, ...]) -> list[str]: """处理 updates 通道:工具事件、待办、子智能体台账。""" from langchain_core.messages import AIMessage, ToolMessage frames: list[str] = [] if not isinstance(data, dict): return frames # 中断会以 __interrupt__ 键出现在 updates 里,比 aget_state 更及时 if "__interrupt__" in data: for item in data.get("__interrupt__") or (): self.interrupt_payloads.append(getattr(item, "value", item)) return frames # 子智能体身份:updates 的 ns 指向它所属的 task 调用 run = self.tracker.note_namespace(ns) if ns else None subagent_type = run.subagent_type if run else "" subagent_id = run.call_id if run else "" for node, update in data.items(): if not isinstance(update, dict): continue todos = update.get("todos") if isinstance(todos, list): self.todos = [t for t in todos if isinstance(t, dict)] frames.append(self.writer.todo_list(self.todos)) for message in update.get("messages") or (): if isinstance(message, ToolMessage): frames.extend(self._on_tool_message(message, subagent_type, subagent_id)) elif isinstance(message, AIMessage): frames.extend(self._on_ai_message(message, subagent_type, subagent_id)) return frames def _on_ai_message( self, message: Any, subagent_type: str, subagent_id: str ) -> list[str]: """处理一条完整的 AIMessage:登记 task 调用、发工具调用卡。""" from tools.registry import get_label frames: list[str] = [] tool_calls = list(getattr(message, "tool_calls", None) or []) # 先让跟踪器登记 task 调用,这样后续子智能体事件能立刻带上身份 for run in self.tracker.note_task_calls(tool_calls): frames.append( self.writer.progress( f"委派子智能体:{run.subagent_type}", subagent=run.subagent_type, subagent_id=run.call_id, description=run.description, ) ) names = [c.get("name", "") for c in tool_calls if isinstance(c, dict)] if names: frames.append( self.writer.executing( f"调用工具:{'、'.join(get_label(n) for n in names)}", cn_agent=self.req.agent_cn_name, tools=names, cn_tools=[get_label(n) for n in names], ) ) for call in tool_calls: if not isinstance(call, dict): continue name = call.get("name") or "" if not name: continue args = call.get("args") self.tool_calls.append({"name": name, "args": args}) frames.append( self.writer.tool_call( name, args=args, cn_tool=get_label(name), agent_id=subagent_id, tool_call_id=str(call.get("id") or ""), ) ) return frames def _on_tool_message( self, message: Any, subagent_type: str, subagent_id: str ) -> list[str]: """处理一条 ToolMessage:发工具结果卡、抽取来源、结子智能体。""" from tools.registry import get_label frames: list[str] = [] name = getattr(message, "name", None) or "tool" content = _message_text(message) # task 的返回即子智能体结束 call_id = getattr(message, "tool_call_id", None) or "" if name == "task" and call_id: run = self.tracker.note_result(call_id, ok=True) if run is not None: frames.append( self.writer.agent_done( run.subagent_type, cn_agent=run.subagent_type, duration_ms=run.duration_ms, progress="已完成", agent_id=run.call_id, ) ) self.meter.note_text(content) self.sources.note_tool_result(name, content) frames.append( self.writer.tool_result( name, content=content[:4000], cn_tool=get_label(name), agent_id=subagent_id, error=str(getattr(message, "status", "") or "") == "error", tool_call_id=call_id, ) ) return frames # -------- messages 通道 -------- def _handle_messages(self, data: Any, ns: tuple[Any, ...]) -> list[str]: """处理 messages 通道:正文 token、思考 token、用量计量。""" from langchain_core.messages import AIMessageChunk frames: list[str] = [] if not isinstance(data, tuple) or len(data) != 2: return frames chunk, metadata = data metadata = metadata if isinstance(metadata, dict) else {} # 摘要中间件的模型调用也在这条通道上:它的 token 不能给用户看 source = metadata.get("lc_source") if source == _SUMMARIZATION_SOURCE: return frames # 用量计量对所有 chunk 都做(包括子智能体的),口径才完整 self.meter.note_message(chunk) # 只有 AIMessageChunk 的内容才是「模型正在说的话」。ToolMessage 也走 # 这条通道,它的 content 是工具返回的原始数据——不加这个判断,工具 # 结果的 JSON 会混进给用户看的正文里。 if not isinstance(chunk, AIMessageChunk): return frames reasoning = _reasoning_text(chunk) if reasoning: frames.append(self.writer.thinking_token(reasoning)) text = _message_text(chunk) if text: self.answer_parts.append(text) frames.append(self.writer.token(text)) return frames # ======================================================== # 中断 # ======================================================== async def _emit_interrupts(self) -> AsyncIterator[str]: """把中断转成 1.5 形状的 interrupt 事件。 两个来源:流里捕获到的 ``__interrupt__``,以及流结束后从图状态里 读出的 pending interrupt。优先用前者(更及时),两者都空则无中断。 """ payloads = list(self.interrupt_payloads) if not payloads: payloads = await self._interrupts_from_state() if not payloads: return for payload in payloads: frame = self._adapt_interrupt(payload) if frame: yield frame async def _interrupts_from_state(self) -> list[Any]: """从图状态里读出挂起的中断载荷。""" get_state = getattr(self.req.agent, "aget_state", None) if get_state is None: return [] try: state = await get_state(self.req.config) except Exception as exc: # pragma: no cover - 取决于底层库 logger.warning("读取图状态失败:%s: %s", type(exc).__name__, exc) return [] found: list[Any] = [] for task in getattr(state, "tasks", None) or (): for item in getattr(task, "interrupts", None) or (): found.append(getattr(item, "value", item)) return found def _adapt_interrupt(self, payload: Any) -> str | None: """官方 HITL 载荷 → 1.5 的 interrupt 事件。 官方形状是 ``{"action_requests":[{name,args,description}], "review_configs":[...]}``; 1.5 前端要的是 ``tool_approval{tool,args,reason,side_effect}``。 这个转换是「协议完全兼容」的落点,改动前先看 sse_events 的契约表。 """ from tools.registry import get_label if not isinstance(payload, dict): return self.writer.interrupt_raw(message=str(payload)[:200]) requests = payload.get("action_requests") if not isinstance(requests, list) or not requests: # 不是工具审批类中断:给一条通用提示,别让前端静默卡住 return self.writer.interrupt_raw( message="智能体暂停,等待输入", payload=payload ) # 官方把同一轮的多个待审批调用合并成一个中断,这里取第一个作为展示主体, # 其余挂在 extra 里——1.5 的审批卡一次只渲染一个工具 first = requests[0] if isinstance(requests[0], dict) else {} tool_name = str(first.get("name") or "") return self.writer.interrupt_tool_approval( tool=tool_name, args=first.get("args"), reason=str(first.get("description") or ""), side_effect="该操作需人工确认后才会执行", cn_tool=get_label(tool_name), pending_count=len(requests), ) # ======================================================== # 收尾 # ======================================================== async def _completion(self, aborted: bool = False) -> AsyncIterator[str]: """落库 → 计费 → 来源 → 上下文用量 → agent_done/done。 无论正常结束、异常还是断连都走这里(幂等,重复调用只生效一次)。 """ if self.finished: return self.finished = True duration_ms = int((time.time() - self.req.started_at) * 1000) answer = "".join(self.answer_parts) if aborted and answer: answer += "\n\n(本次回复被中断,内容可能不完整)" # 1) 上下文用量快照。**只取一次**:snapshot() 会消费掉压缩存档的待办标记, # 取两次会导致第二次拿不到存档内容。 self._apply_estimate_fallback() snapshot = self.meter.snapshot() # 2) 计费:把上下文计量器的数字转给计费计量器。 # 这是两个计量器唯一的连接点——漏了这一步,交互式对话全程不会计费, # 而 /usage/today 会显示 0,表面上「一切正常」。 self.billing.add(snapshot.input_tokens, snapshot.output_tokens, snapshot.usage_source) # 3) 落库:assistant 消息是会话历史的事实源 citations = self.sources.build_citations() if answer or citations: try: chat_store.save_message( self.req.session_id, "assistant", answer, tool_calls=self.tool_calls or None, citations=citations or None, ) except Exception as exc: logger.error("落库失败 session=%s:%s: %s", self.req.session_id, type(exc).__name__, exc) # 4) 上下文用量 + 压缩存档 await self._persist_context(snapshot) # 5) 计费入账(幂等) self.billing.settle() # 6) 来源标注 yield self.writer.citations(citations, model_only=not citations) if aborted: yield self.writer.done(duration_ms=duration_ms, aborted=True) return yield self.writer.agent_done( self.req.agent_name, cn_agent=self.req.agent_cn_name, duration_ms=duration_ms, progress="已完成", **self.meter.summary_fields(), ) yield self.writer.done(duration_ms=duration_ms, usage=self._usage_payload()) def _usage_payload(self) -> dict[str, Any]: """给前端的用量摘要(便于测试页显示,1.5 前端会忽略未知字段)。""" return { "input_tokens": self.meter.input_tokens, "output_tokens": self.meter.output_tokens, "total_tokens": self.meter.total_tokens, "usage_source": self.meter.usage_source, "context_tokens": self.meter.context_tokens, } def _apply_estimate_fallback(self) -> None: """provider 没返回任何 usage 时,用缓冲区文本估算兜底。 不兜底的话,这类请求会以 0 token 入账——用户白用、额度不受控, 而且从账面上完全看不出来。兜底后至少 usage_source=estimate 能 标示这批数字不可尽信。 """ if self.meter.total_tokens > 0: return estimated = self.meter.estimate_from_buffer() if estimated > 0: logger.warning( "本次未取到 provider usage,按缓冲区估算 %d tokens(usage_source=estimate)" " session=%s", estimated, self.req.session_id, ) self.meter.note_estimated(estimated) async def _persist_context(self, snapshot: Any) -> None: """落 session_context_usage,并在发生压缩时登记原文存档位置。 存档优先用**框架自己的产物**:摘要中间件会把被淘汰的历史追加写到 ``/conversation_history/{thread_id}.md``,并把路径嵌进摘要正文—— 那份文件是权威的,agent 也能自己读回来。只有框架产物不存在时才 回落写我们自己的缓冲区,避免「以为有存档其实没有」。 """ archive_path = "" try: if snapshot.compacted: archive_path = await asyncio.to_thread( self._archive_compaction, snapshot.archive_text, snapshot.archive_tokens ) chat_store.save_context_usage( session_id=self.req.session_id, turn=chat_store.count_messages(self.req.session_id) // 2, input_tokens=snapshot.input_tokens, output_tokens=snapshot.output_tokens, cache_read_tokens=snapshot.cache_read_tokens, usage_source=snapshot.usage_source, compacted=snapshot.compacted, compact_count=snapshot.compact_count, snapshot={ "context_tokens": snapshot.context_tokens, "context_window": self.meter.window, "model_calls": snapshot.model_calls, "archive_path": archive_path, }, ) except Exception as exc: logger.error( "上下文用量落库失败 session=%s:%s: %s", self.req.session_id, type(exc).__name__, exc, ) def _archive_compaction(self, buffered_text: str | None, buffered_tokens: int) -> str: """登记(必要时写出)压缩前原文,返回存档路径或空串。""" history_dir = config.ROOT / "conversation_history" framework_file = history_dir / f"{self.req.thread_id}.md" if framework_file.exists(): path = str(framework_file) chat_store.save_context_archive( self.req.session_id, 0, path, buffered_tokens ) logger.info("压缩原文由框架落盘:%s", path) return path if not buffered_text: logger.warning( "检测到压缩但既无框架落盘也无本地缓冲,无法存档 session=%s", self.req.session_id, ) return "" archive_dir = config.ROOT / "data" / "context_archives" archive_dir.mkdir(parents=True, exist_ok=True) stamp = time.strftime("%Y%m%d-%H%M%S") target = archive_dir / f"{self.req.session_id}-{stamp}.md" target.write_text(buffered_text, encoding="utf-8") chat_store.save_context_archive( self.req.session_id, 0, str(target), buffered_tokens ) logger.info("压缩原文由本进程落盘:%s(%d tokens)", target, buffered_tokens) return str(target) # ======================================================== # 异常 # ======================================================== async def _error_frames(self, exc: Exception) -> AsyncIterator[str]: """把异常翻译成 error 帧 + 收尾帧。""" from langgraph.errors import GraphRecursionError if isinstance(exc, GraphRecursionError): logger.warning("图递归上限触发 session=%s:%s", self.req.session_id, exc) yield self.writer.error( "任务步骤过多已被中止,请把任务拆成更小的步骤重试", code="recursion_limit" ) async for frame in self._completion(aborted=True): yield frame return if isinstance(exc, concurrency.TaskWallTimeoutError): logger.warning("任务墙钟超时 session=%s:%s", self.req.session_id, exc) yield self.writer.error(str(exc), code="task_timeout") async for frame in self._completion(aborted=True): yield frame return logger.exception("流式执行异常 session=%s", self.req.session_id) yield self.writer.error( f"服务内部错误:{type(exc).__name__}: {exc}", code="internal_error" ) async for frame in self._completion(aborted=True): yield frame # ======================================================== # 资源 # ======================================================== def _finalize(self) -> None: """释放并发位并兜底结算。幂等。""" try: if not self.billing.settled: self.billing.settle() finally: if self.req.slot is not None: self.req.slot.release() def _meter_for_request() -> ContextMeter: """给本轮请求建一个上下文计量器,窗口取自模型配置。""" from agents import model_factory try: window = model_factory.context_window() except Exception: # pragma: no cover - 配置异常不该阻断对话 window = 0 return ContextMeter(window=window) def sse_headers(session_id: str) -> dict[str, str]: """SSE 响应头。``X-Accel-Buffering: no`` 是给 nginx 的:不关掉的话 反代会缓冲整个响应,流式效果全无(表现为「一次性吐出全部内容」)。""" return { "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", "X-Session-Id": session_id, } async def pending_interrupts(agent: Any, config: dict[str, Any]) -> list[Any]: """读取当前挂起的中断载荷列表。 供两处复用:``GET /chat/interrupt``(刷新页面后补渲染卡片)与 ``POST /chat/resume``(在构造 decisions 前先知道有几个待审批动作—— 官方 HITL 要求 decisions 与 action_requests **数量严格一一对应**, 数量对不上会直接抛错)。 """ get_state = getattr(agent, "aget_state", None) if get_state is None: return [] try: state = await get_state(config) except Exception as exc: logger.warning("读取图状态失败:%s: %s", type(exc).__name__, exc) return [] found: list[Any] = [] for task in getattr(state, "tasks", None) or (): for item in getattr(task, "interrupts", None) or (): found.append(getattr(item, "value", item)) return found