sse_core.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819
  1. """SSE 流式内核。
  2. 职责:把 LangGraph 的 ``astream`` 事件流翻译成 1.5 兼容的 SSE 帧,并在
  3. 结束时完成落库、计费、来源标注与上下文计量。
  4. ## 通道分工(重要,避免重复事件)
  5. ``stream_mode=["updates", "messages"]`` 两条通道都在推同一个 agent 的进展,
  6. 不加区分会导致同一个工具调用被发两次。这里的分工是:
  7. - **updates 通道**:工具调用/结果、待办清单、子智能体生命周期、状态提示。
  8. 走这里的原因是它给的是**完整的消息对象**(``tool_calls`` 里 args 已解析
  9. 完毕、``ToolMessage`` 已成对),不需要在流式增量里拼半个 JSON。
  10. - **messages 通道**:正文 token、思考 token、真实 token 用量。
  11. ``version="v2"`` 下每个 chunk 是 ``{"type", "ns", "data"}``;``subgraphs=True``
  12. 让子智能体的产出带 ``ns``。两个格式都做了兼容(旧版是 ``(mode, data)`` 元组)。
  13. ## 与 1.5 的差异
  14. 1.5 的 sse_core 是 1908 行,里面混着计费、压缩、Word 导出、来源刮取等多条
  15. 业务线。2.0 把这些拆成独立组件(context_meter / source_tracker /
  16. subagent_tracker),内核只负责「读流 → 派事件 → 收尾」。
  17. """
  18. from __future__ import annotations
  19. import asyncio
  20. import time
  21. from dataclasses import dataclass, field
  22. from pathlib import Path
  23. from datetime import datetime
  24. import anyio
  25. from typing import Any, AsyncIterator
  26. import config
  27. from api import concurrency
  28. from api.interrupts import pack, interrupt_frame
  29. from api.context_meter import ContextMeter
  30. from api.sse_events import EventWriter
  31. from api.source_tracker import SourceTracker
  32. from api.subagent_tracker import SubagentTracker
  33. from billing.meter import UsageMeter
  34. from db import chat_store, db_conn, run_store
  35. from obs import get_logger, get_trace_id
  36. logger = get_logger(__name__)
  37. MAIN_AGENT_NAME = "deep-agent"
  38. MAIN_AGENT_CN = "深度智能体"
  39. # 摘要中间件的模型调用也在 messages 通道里推 token,靠这个标记过滤掉,
  40. # 否则摘要正文会混进给用户看的回答里
  41. _SUMMARIZATION_SOURCE = "summarization"
  42. @dataclass
  43. class StreamRequest:
  44. """一次流式请求的全部入参。"""
  45. agent: Any
  46. session_id: str
  47. thread_id: str
  48. config: dict[str, Any]
  49. message: str
  50. user_name: str = "default"
  51. kind: str = "chat"
  52. model: str = ""
  53. original_user_message: str = ""
  54. trace_id: str = ""
  55. slot: concurrency.TaskSlot | None = None
  56. started_at: float = field(default_factory=time.time)
  57. agent_name: str = MAIN_AGENT_NAME
  58. agent_cn_name: str = MAIN_AGENT_CN
  59. # 非 None 表示这是一次中断恢复:不组装消息,直接把 Command(resume=...) 喂回图
  60. resume_payload: Any = None
  61. session_lease: Any = None
  62. run_id: str = field(default_factory=db_conn.new_id)
  63. def normalize_chunk(chunk: Any) -> tuple[str | None, tuple[Any, ...], Any]:
  64. """把不同版本的 astream chunk 归一成 ``(mode, ns, data)``。
  65. 兼容三种形态:v2 的 ``{"type","ns","data"}`` 字典、旧版 ``(mode, data)``
  66. 元组、以及无法识别的原样返回(mode=None,上层据此忽略)。
  67. """
  68. if isinstance(chunk, dict) and "data" in chunk and "type" in chunk:
  69. return str(chunk["type"]), tuple(chunk.get("ns") or ()), chunk["data"]
  70. if isinstance(chunk, tuple) and len(chunk) == 2 and isinstance(chunk[0], str):
  71. return chunk[0], (), chunk[1]
  72. return None, (), chunk
  73. def _message_text(message: Any) -> str:
  74. """把消息 content 拍平成纯文本(content 可能是字符串或内容块列表)。"""
  75. content = getattr(message, "content", None)
  76. if isinstance(content, str):
  77. return content
  78. if isinstance(content, list):
  79. parts: list[str] = []
  80. for block in content:
  81. if isinstance(block, str):
  82. parts.append(block)
  83. elif isinstance(block, dict):
  84. text = block.get("text")
  85. if isinstance(text, str):
  86. parts.append(text)
  87. return "".join(parts)
  88. return ""
  89. def _reasoning_text(message: Any) -> str:
  90. """抽取思考链文本(deepseek 的 reasoning_content)。"""
  91. extra = getattr(message, "additional_kwargs", None)
  92. if isinstance(extra, dict):
  93. value = extra.get("reasoning_content")
  94. if isinstance(value, str):
  95. return value
  96. return ""
  97. class ChatStream:
  98. """一次 /chat 请求的流式执行。"""
  99. def __init__(self, request: StreamRequest) -> None:
  100. self.req = request
  101. self.writer = EventWriter(request.session_id, request.trace_id or get_trace_id())
  102. self.meter = _meter_for_request()
  103. self.tracker = SubagentTracker()
  104. self.sources = SourceTracker()
  105. self.billing = UsageMeter(
  106. user_name=request.user_name,
  107. session_id=request.session_id,
  108. kind=request.kind,
  109. model=request.model or config.MODEL,
  110. trace_id=request.trace_id,
  111. usage_id=request.run_id,
  112. usage_day=db_conn.day_key(datetime.fromtimestamp(request.started_at)),
  113. )
  114. # 累计状态
  115. self.answer_parts: list[str] = []
  116. self.tool_calls: list[dict[str, Any]] = []
  117. self.todos: list[dict[str, Any]] = []
  118. self.interrupt_payloads: list[Any] = []
  119. self.finished = False
  120. self.status = "running"
  121. self.error = ""
  122. self._billing_added = False
  123. self._started = False
  124. self._snapshot = None
  125. self._answer_saved = False
  126. self._settlement_task: asyncio.Task | None = None
  127. # ========================================================
  128. # 主入口
  129. # ========================================================
  130. async def generate(self) -> AsyncIterator[str]:
  131. """产出整个 SSE 流。保证:计费幂等入账、并发位释放、图游标关闭。
  132. 这里显式绑定 trace_id 而不是依赖 HTTP 中间件设的 ContextVar:
  133. 响应体的迭代在 Starlette 的中间件调用栈之外执行,中间件里绑的值
  134. 在生成器里不保证可见。绑定发生在生成器内部,生命周期与本次流一致。
  135. """
  136. from obs import bind_trace_id, unbind_trace_id
  137. token = bind_trace_id(self.req.trace_id or None)
  138. stream = None
  139. try:
  140. await asyncio.to_thread(run_store.start, self.req.run_id, self.req.session_id,
  141. self.req.user_name, self.req.kind)
  142. self._started = True
  143. yield self.writer.agent_start(self.req.agent_name, self.req.agent_cn_name)
  144. yield self.writer.thinking("正在理解您的问题...", self.req.agent_cn_name)
  145. graph_input = await self._prepare_input()
  146. stream = self.req.agent.astream(
  147. graph_input,
  148. stream_mode=["updates", "messages"],
  149. config=self.req.config,
  150. subgraphs=True,
  151. version="v2",
  152. )
  153. async for frame in self._drive(stream):
  154. yield frame
  155. # 中断检测必须在流结束后做:interrupt 会让流自然结束
  156. async for frame in self._emit_interrupts():
  157. yield frame
  158. async for frame in self._completion():
  159. yield frame
  160. except (asyncio.CancelledError, GeneratorExit):
  161. self.status = "cancelled"
  162. raise
  163. except Exception as exc:
  164. self.status = "failed"
  165. self.error = type(exc).__name__
  166. async for frame in self._error_frames(exc):
  167. yield frame
  168. finally:
  169. # 清理与发帧独立。ASGI 断连的取消作用域不能打断保存和释放。
  170. with anyio.CancelScope(shield=True):
  171. cleanup = asyncio.create_task(self._cleanup(stream))
  172. try:
  173. await asyncio.shield(cleanup)
  174. except asyncio.CancelledError:
  175. await cleanup
  176. unbind_trace_id(token)
  177. async def _cleanup(self, stream: Any) -> None:
  178. try:
  179. if stream is not None:
  180. close = getattr(stream, "aclose", None)
  181. if close is not None:
  182. try:
  183. await close()
  184. except Exception:
  185. logger.exception("关闭图事件流失败 session=%s", self.req.session_id)
  186. await self._settle(aborted=self.status in {"cancelled", "failed"})
  187. finally:
  188. try:
  189. await asyncio.to_thread(self._finalize)
  190. finally:
  191. if self.req.session_lease is not None:
  192. self.req.session_lease.release()
  193. # ========================================================
  194. # 输入准备
  195. # ========================================================
  196. async def _prepare_input(self) -> Any:
  197. """组装喂给图的输入。
  198. 两种情形:
  199. - **恢复**(``resume_payload`` 非 None):喂 ``Command(resume=...)``,
  200. 不组装消息——图会从挂起的那个节点继续。
  201. - **新请求**:热线程(checkpoint 里已有状态)只喂当轮新消息;冷线程
  202. 需要从会话历史补种。
  203. 补种这一步不能省,也不能对所有线程都做:热线程若把历史重复追加,
  204. 上下文会随轮次按 O(N²) 膨胀——这是长会话最容易被忽略的膨胀来源。
  205. """
  206. if self.req.resume_payload is not None:
  207. from langgraph.types import Command
  208. logger.info("恢复中断 session=%s", self.req.session_id)
  209. return Command(resume=self.req.resume_payload)
  210. from langchain_core.messages import AIMessage, HumanMessage
  211. from agents import agent_memory
  212. messages: list[Any] = []
  213. warm = await agent_memory.has_checkpoint(self.req.thread_id)
  214. if not warm:
  215. history = await asyncio.to_thread(chat_store.get_messages,
  216. self.req.session_id, limit=config.CTX_RECENT_LIMIT
  217. )
  218. for row in history:
  219. content = row.get("content") or ""
  220. if not content:
  221. continue
  222. if row.get("msg_role") == "user":
  223. messages.append(HumanMessage(content=content))
  224. else:
  225. messages.append(AIMessage(content=content))
  226. if messages:
  227. logger.info(
  228. "冷线程补种历史 session=%s 条数=%d", self.req.session_id, len(messages)
  229. )
  230. else:
  231. logger.info("热线程,仅喂当轮消息 session=%s", self.req.session_id)
  232. # 用户消息先落库:即使这一轮流到一半挂了,历史里也得有用户问了什么
  233. await asyncio.to_thread(chat_store.save_message, self.req.session_id, "user",
  234. self.req.original_user_message or self.req.message)
  235. self.meter.note_text(self.req.message)
  236. messages.append(HumanMessage(content=self.req.message))
  237. # 必须是 {"messages": [...]} 字典:直接传 list 会被当成 __start__ 节点的
  238. # 返回值,图会报 InvalidUpdateError: Expected dict
  239. return {"messages": messages}
  240. # ========================================================
  241. # 流驱动
  242. # ========================================================
  243. async def _drive(self, stream: Any) -> AsyncIterator[str]:
  244. """消费 astream,派发事件帧。带墙钟上限。"""
  245. remaining = config.TASK_WALL_TIMEOUT - (time.time() - self.req.started_at)
  246. deadline = time.monotonic() + max(remaining, 0)
  247. timed_out = False
  248. iterator = stream.__aiter__()
  249. while True:
  250. if time.monotonic() >= deadline:
  251. timed_out = True
  252. break
  253. try:
  254. chunk = await asyncio.wait_for(
  255. iterator.__anext__(), timeout=max(deadline - time.monotonic(), 0.001)
  256. )
  257. except StopAsyncIteration:
  258. break
  259. except asyncio.TimeoutError:
  260. timed_out = True
  261. break
  262. mode, ns, data = normalize_chunk(chunk)
  263. if mode == "updates":
  264. for frame in self._handle_updates(data, ns):
  265. yield frame
  266. elif mode == "messages":
  267. for frame in self._handle_messages(data, ns):
  268. yield frame
  269. # mode 为 None:无法识别的 chunk,忽略而不是猜
  270. if timed_out:
  271. raise concurrency.TaskWallTimeoutError(
  272. f"任务超过墙钟上限 {config.TASK_WALL_TIMEOUT}s 已被中止"
  273. )
  274. # 子智能体可能因为异常/中断没收到结果,收尾时统一关掉,避免前端卡片长转
  275. for run in self.tracker.finish_all():
  276. yield self.writer.progress(
  277. f"子智能体 {run.subagent_type} 未正常返回,已中止",
  278. subagent=run.subagent_type, subagent_id=run.call_id,
  279. )
  280. # -------- updates 通道 --------
  281. def _handle_updates(self, data: Any, ns: tuple[Any, ...]) -> list[str]:
  282. """处理 updates 通道:工具事件、待办、子智能体台账。"""
  283. from langchain_core.messages import AIMessage, ToolMessage
  284. frames: list[str] = []
  285. if not isinstance(data, dict):
  286. return frames
  287. # 中断会以 __interrupt__ 键出现在 updates 里,比 aget_state 更及时
  288. if "__interrupt__" in data:
  289. for item in data.get("__interrupt__") or ():
  290. self.interrupt_payloads.append(pack(item))
  291. return frames
  292. # 子智能体身份:updates 的 ns 指向它所属的 task 调用
  293. run = self.tracker.note_namespace(ns) if ns else None
  294. subagent_type = run.subagent_type if run else ""
  295. subagent_id = run.call_id if run else ""
  296. for node, update in data.items():
  297. if not isinstance(update, dict):
  298. continue
  299. todos = update.get("todos")
  300. if isinstance(todos, list):
  301. self.todos = [t for t in todos if isinstance(t, dict)]
  302. frames.append(self.writer.todo_list(self.todos))
  303. for message in update.get("messages") or ():
  304. if isinstance(message, ToolMessage):
  305. frames.extend(self._on_tool_message(message, subagent_type, subagent_id))
  306. elif isinstance(message, AIMessage):
  307. frames.extend(self._on_ai_message(message, subagent_type, subagent_id))
  308. return frames
  309. def _on_ai_message(
  310. self, message: Any, subagent_type: str, subagent_id: str
  311. ) -> list[str]:
  312. """处理一条完整的 AIMessage:登记 task 调用、发工具调用卡。"""
  313. from tools.registry import get_label
  314. frames: list[str] = []
  315. tool_calls = list(getattr(message, "tool_calls", None) or [])
  316. # 先让跟踪器登记 task 调用,这样后续子智能体事件能立刻带上身份
  317. for run in self.tracker.note_task_calls(tool_calls):
  318. frames.append(
  319. self.writer.progress(
  320. f"委派子智能体:{run.subagent_type}",
  321. subagent=run.subagent_type,
  322. subagent_id=run.call_id,
  323. description=run.description,
  324. )
  325. )
  326. names = [c.get("name", "") for c in tool_calls if isinstance(c, dict)]
  327. if names:
  328. frames.append(
  329. self.writer.executing(
  330. f"调用工具:{'、'.join(get_label(n) for n in names)}",
  331. cn_agent=self.req.agent_cn_name,
  332. tools=names,
  333. cn_tools=[get_label(n) for n in names],
  334. )
  335. )
  336. for call in tool_calls:
  337. if not isinstance(call, dict):
  338. continue
  339. name = call.get("name") or ""
  340. if not name:
  341. continue
  342. args = call.get("args")
  343. self.tool_calls.append({"name": name, "args": args})
  344. frames.append(
  345. self.writer.tool_call(
  346. name,
  347. args=args,
  348. cn_tool=get_label(name),
  349. agent_id=subagent_id,
  350. tool_call_id=str(call.get("id") or ""),
  351. )
  352. )
  353. return frames
  354. def _on_tool_message(
  355. self, message: Any, subagent_type: str, subagent_id: str
  356. ) -> list[str]:
  357. """处理一条 ToolMessage:发工具结果卡、抽取来源、结子智能体。"""
  358. from tools.registry import get_label
  359. frames: list[str] = []
  360. name = getattr(message, "name", None) or "tool"
  361. content = _message_text(message)
  362. # task 的返回即子智能体结束
  363. call_id = getattr(message, "tool_call_id", None) or ""
  364. if name == "task" and call_id:
  365. run = self.tracker.note_result(call_id, ok=True)
  366. if run is not None:
  367. frames.append(
  368. self.writer.agent_done(
  369. run.subagent_type,
  370. cn_agent=run.subagent_type,
  371. duration_ms=run.duration_ms,
  372. progress="已完成",
  373. agent_id=run.call_id,
  374. )
  375. )
  376. self.meter.note_text(content)
  377. self.sources.note_tool_result(name, content)
  378. frames.append(
  379. self.writer.tool_result(
  380. name,
  381. content=content[:4000],
  382. cn_tool=get_label(name),
  383. agent_id=subagent_id,
  384. error=str(getattr(message, "status", "") or "") == "error",
  385. tool_call_id=call_id,
  386. )
  387. )
  388. return frames
  389. # -------- messages 通道 --------
  390. def _handle_messages(self, data: Any, ns: tuple[Any, ...]) -> list[str]:
  391. """处理 messages 通道:正文 token、思考 token、用量计量。"""
  392. from langchain_core.messages import AIMessage, AIMessageChunk
  393. frames: list[str] = []
  394. if not isinstance(data, tuple) or len(data) != 2:
  395. return frames
  396. chunk, metadata = data
  397. metadata = metadata if isinstance(metadata, dict) else {}
  398. # 摘要中间件的模型调用也在这条通道上:它的 token 不能给用户看
  399. source = metadata.get("lc_source")
  400. # 汇总全部模型费用;只有主线程更新当前上下文窗口占用。
  401. self.meter.note_message(chunk, is_main_context=not ns and source != _SUMMARIZATION_SOURCE)
  402. if source == _SUMMARIZATION_SOURCE or ns:
  403. return frames
  404. # 只有 AIMessageChunk 的内容才是「模型正在说的话」。ToolMessage 也走
  405. # 这条通道,它的 content 是工具返回的原始数据——不加这个判断,工具
  406. # 结果的 JSON 会混进给用户看的正文里。
  407. if not isinstance(chunk, (AIMessage, AIMessageChunk)):
  408. return frames
  409. reasoning = _reasoning_text(chunk)
  410. if reasoning:
  411. frames.append(self.writer.thinking_token(reasoning))
  412. text = _message_text(chunk)
  413. if text:
  414. self.meter.note_text(text)
  415. self.answer_parts.append(text)
  416. frames.append(self.writer.token(text))
  417. return frames
  418. # ========================================================
  419. # 中断
  420. # ========================================================
  421. async def _emit_interrupts(self) -> AsyncIterator[str]:
  422. """把中断转成 1.5 形状的 interrupt 事件。
  423. 两个来源:流里捕获到的 ``__interrupt__``,以及流结束后从图状态里
  424. 读出的 pending interrupt。优先用前者(更及时),两者都空则无中断。
  425. """
  426. payloads = list(self.interrupt_payloads)
  427. if not payloads:
  428. payloads = await self._interrupts_from_state()
  429. if not payloads:
  430. return
  431. self.interrupt_payloads = payloads
  432. for payload in payloads:
  433. frame = self._adapt_interrupt(payload)
  434. if frame:
  435. yield frame
  436. async def _interrupts_from_state(self) -> list[Any]:
  437. """从图状态里读出挂起的中断载荷。"""
  438. get_state = getattr(self.req.agent, "aget_state", None)
  439. if get_state is None:
  440. return []
  441. try:
  442. state = await get_state(self.req.config)
  443. except Exception as exc: # pragma: no cover - 取决于底层库
  444. logger.warning("读取图状态失败:%s: %s", type(exc).__name__, exc)
  445. return []
  446. found: list[Any] = []
  447. for task in getattr(state, "tasks", None) or ():
  448. for item in getattr(task, "interrupts", None) or ():
  449. found.append(pack(item))
  450. return found
  451. def _adapt_interrupt(self, payload: Any) -> str | None:
  452. return interrupt_frame(self.writer, payload)
  453. # ========================================================
  454. # 收尾
  455. # ========================================================
  456. async def _settle(self, aborted: bool = False) -> None:
  457. # 断连恰逢数据库写入时继续等待同一次收尾,避免在线程仍写入时重做。
  458. if self._settlement_task is None:
  459. self._settlement_task = asyncio.create_task(self._settle_once(aborted))
  460. await asyncio.shield(self._settlement_task)
  461. async def _settle_once(self, aborted: bool = False) -> None:
  462. """落库 → 计费 → 来源 → 上下文用量 → agent_done/done。
  463. 无论正常结束、异常还是断连都走这里(幂等,重复调用只生效一次)。
  464. """
  465. if self.finished:
  466. return
  467. duration_ms = int((time.time() - self.req.started_at) * 1000)
  468. answer = "".join(self.answer_parts)
  469. if aborted and answer:
  470. answer += "\n\n(本次回复被中断,内容可能不完整)"
  471. # 1) 上下文用量快照。**只取一次**:snapshot() 会消费掉压缩存档的待办标记,
  472. # 取两次会导致第二次拿不到存档内容。
  473. self._apply_estimate_fallback()
  474. if self._snapshot is None:
  475. self._snapshot = self.meter.snapshot()
  476. snapshot = self._snapshot
  477. # 2) 计费:把上下文计量器的数字转给计费计量器。
  478. # 这是两个计量器唯一的连接点——漏了这一步,交互式对话全程不会计费,
  479. # 而 /usage/today 会显示 0,表面上「一切正常」。
  480. if not self._billing_added:
  481. self.billing.add(snapshot.input_tokens, snapshot.output_tokens, snapshot.usage_source)
  482. self._billing_added = True
  483. # 3) 落库:assistant 消息是会话历史的事实源
  484. citations = self.sources.build_citations()
  485. if (answer or citations) and not self._answer_saved:
  486. try:
  487. await asyncio.to_thread(
  488. chat_store.save_message, self.req.session_id,
  489. "assistant",
  490. answer,
  491. tool_calls=self.tool_calls or None,
  492. citations=citations or None,
  493. )
  494. self._answer_saved = True
  495. except Exception as exc:
  496. self.status = "failed"
  497. self.error = "message_persistence_failed"
  498. logger.error("落库失败 session=%s:%s: %s", self.req.session_id, type(exc).__name__, exc)
  499. # 4) 上下文用量 + 压缩存档
  500. await self._persist_context(snapshot)
  501. # 5) 计费入账(幂等)
  502. await asyncio.to_thread(self.billing.settle)
  503. if self.status == "running":
  504. self.status = "waiting_input" if self.interrupt_payloads else "success"
  505. if self._started:
  506. await asyncio.to_thread(run_store.finish, self.req.run_id, self.status, self.error)
  507. # 用户在普通恢复接口完成定时任务审批后,同步其执行记录。
  508. if self.req.kind == "resume":
  509. await asyncio.to_thread(db_conn.execute,
  510. "UPDATE schedule_runs SET status = ?, finished_at = ? WHERE thread_id = ? AND status = 'waiting_input'",
  511. (self.status, None if self.status == "waiting_input" else db_conn.fmt_dt(), self.req.session_id))
  512. self.finished = True
  513. async def _completion(self, aborted: bool = False) -> AsyncIterator[str]:
  514. await self._settle(aborted)
  515. duration_ms = int((time.time() - self.req.started_at) * 1000)
  516. citations = self.sources.build_citations()
  517. # 6) 来源标注
  518. yield self.writer.citations(citations, model_only=not citations)
  519. if aborted or self.status in {"failed", "cancelled"}:
  520. if self.error == "message_persistence_failed":
  521. yield self.writer.error("回复未能保存,请稍后检查历史记录", code=self.error)
  522. yield self.writer.done(duration_ms=duration_ms, aborted=True)
  523. return
  524. if self.status == "waiting_input":
  525. yield self.writer.done(duration_ms=duration_ms, waiting_input=True, run_id=self.req.run_id, usage=self._usage_payload())
  526. return
  527. yield self.writer.agent_done(
  528. self.req.agent_name,
  529. cn_agent=self.req.agent_cn_name,
  530. duration_ms=duration_ms,
  531. progress="已完成",
  532. **self.meter.summary_fields(),
  533. )
  534. yield self.writer.done(duration_ms=duration_ms, run_id=self.req.run_id, usage=self._usage_payload())
  535. def _usage_payload(self) -> dict[str, Any]:
  536. """给前端的用量摘要(便于测试页显示,1.5 前端会忽略未知字段)。"""
  537. return {
  538. "input_tokens": self.meter.input_tokens,
  539. "output_tokens": self.meter.output_tokens,
  540. "total_tokens": self.meter.total_tokens,
  541. "usage_source": self.meter.usage_source,
  542. "context_tokens": self.meter.context_tokens,
  543. }
  544. def _apply_estimate_fallback(self) -> None:
  545. """provider 没返回任何 usage 时,用缓冲区文本估算兜底。
  546. 不兜底的话,这类请求会以 0 token 入账——用户白用、额度不受控,
  547. 而且从账面上完全看不出来。兜底后至少 usage_source=estimate 能
  548. 标示这批数字不可尽信。
  549. """
  550. if self.meter.total_tokens > 0:
  551. return
  552. estimated = self.meter.estimate_from_buffer()
  553. if estimated > 0:
  554. logger.warning(
  555. "本次未取到 provider usage,按缓冲区估算 %d tokens(usage_source=estimate)"
  556. " session=%s",
  557. estimated, self.req.session_id,
  558. )
  559. self.meter.note_estimated(estimated)
  560. async def _persist_context(self, snapshot: Any) -> None:
  561. """落 session_context_usage,并在发生压缩时登记原文存档位置。
  562. 存档优先用**框架自己的产物**:摘要中间件会把被淘汰的历史追加写到
  563. ``/conversation_history/{thread_id}.md``,并把路径嵌进摘要正文——
  564. 那份文件是权威的,agent 也能自己读回来。只有框架产物不存在时才
  565. 回落写我们自己的缓冲区,避免「以为有存档其实没有」。
  566. """
  567. archive_path = ""
  568. try:
  569. if snapshot.compacted:
  570. archive_path = await asyncio.to_thread(
  571. self._archive_compaction, snapshot.archive_text, snapshot.archive_tokens
  572. )
  573. await asyncio.to_thread(chat_store.save_context_usage,
  574. session_id=self.req.session_id,
  575. turn=chat_store.count_messages(self.req.session_id) // 2,
  576. input_tokens=snapshot.input_tokens,
  577. output_tokens=snapshot.output_tokens,
  578. cache_read_tokens=snapshot.cache_read_tokens,
  579. usage_source=snapshot.usage_source,
  580. compacted=snapshot.compacted,
  581. compact_count=snapshot.compact_count,
  582. snapshot={
  583. "context_tokens": snapshot.context_tokens,
  584. "context_window": self.meter.window,
  585. "model_calls": snapshot.model_calls,
  586. "archive_path": archive_path,
  587. },
  588. )
  589. except Exception as exc:
  590. logger.error(
  591. "上下文用量落库失败 session=%s:%s: %s",
  592. self.req.session_id, type(exc).__name__, exc,
  593. )
  594. def _archive_compaction(self, buffered_text: str | None, buffered_tokens: int) -> str:
  595. """登记(必要时写出)压缩前原文,返回存档路径或空串。"""
  596. from agents.workspace import session_root
  597. history_dir = session_root(self.req.user_name, self.req.session_id) / "conversation_history"
  598. framework_file = history_dir / f"{self.req.thread_id}.md"
  599. if framework_file.exists():
  600. path = str(framework_file)
  601. chat_store.save_context_archive(
  602. self.req.session_id, 0, path, buffered_tokens
  603. )
  604. logger.info("压缩原文由框架落盘:%s", path)
  605. return path
  606. if not buffered_text:
  607. logger.warning(
  608. "检测到压缩但既无框架落盘也无本地缓冲,无法存档 session=%s",
  609. self.req.session_id,
  610. )
  611. return ""
  612. archive_dir = history_dir
  613. archive_dir.mkdir(parents=True, exist_ok=True)
  614. stamp = time.strftime("%Y%m%d-%H%M%S")
  615. target = archive_dir / f"{self.req.session_id}-{stamp}.md"
  616. target.write_text(buffered_text, encoding="utf-8")
  617. chat_store.save_context_archive(
  618. self.req.session_id, 0, str(target), buffered_tokens
  619. )
  620. logger.info("压缩原文由本进程落盘:%s(%d tokens)", target, buffered_tokens)
  621. return str(target)
  622. # ========================================================
  623. # 异常
  624. # ========================================================
  625. async def _error_frames(self, exc: Exception) -> AsyncIterator[str]:
  626. """把异常翻译成 error 帧 + 收尾帧。"""
  627. from langgraph.errors import GraphRecursionError
  628. if isinstance(exc, GraphRecursionError):
  629. logger.warning("图递归上限触发 session=%s:%s", self.req.session_id, exc)
  630. yield self.writer.error(
  631. "任务步骤过多已被中止,请把任务拆成更小的步骤重试", code="recursion_limit"
  632. )
  633. async for frame in self._completion(aborted=True):
  634. yield frame
  635. return
  636. if isinstance(exc, concurrency.TaskWallTimeoutError):
  637. logger.warning("任务墙钟超时 session=%s:%s", self.req.session_id, exc)
  638. yield self.writer.error(str(exc), code="task_timeout")
  639. async for frame in self._completion(aborted=True):
  640. yield frame
  641. return
  642. logger.exception("流式执行异常 session=%s", self.req.session_id)
  643. yield self.writer.error(
  644. f"服务内部错误:{type(exc).__name__}", code="internal_error"
  645. )
  646. async for frame in self._completion(aborted=True):
  647. yield frame
  648. # ========================================================
  649. # 资源
  650. # ========================================================
  651. def _finalize(self) -> None:
  652. """释放并发位并兜底结算。幂等。"""
  653. try:
  654. if not self.billing.settled:
  655. self.billing.settle()
  656. finally:
  657. if self.req.slot is not None:
  658. self.req.slot.release()
  659. def _meter_for_request() -> ContextMeter:
  660. """给本轮请求建一个上下文计量器,窗口取自模型配置。"""
  661. from agents import model_factory
  662. try:
  663. window = model_factory.context_window()
  664. except Exception: # pragma: no cover - 配置异常不该阻断对话
  665. window = 0
  666. return ContextMeter(window=window)
  667. def sse_headers(session_id: str) -> dict[str, str]:
  668. """SSE 响应头。``X-Accel-Buffering: no`` 是给 nginx 的:不关掉的话
  669. 反代会缓冲整个响应,流式效果全无(表现为「一次性吐出全部内容」)。"""
  670. return {
  671. "Cache-Control": "no-cache",
  672. "Connection": "keep-alive",
  673. "X-Accel-Buffering": "no",
  674. "X-Session-Id": session_id,
  675. }
  676. async def pending_interrupts(agent: Any, config: dict[str, Any]) -> list[Any]:
  677. """读取当前挂起的中断载荷列表。
  678. 供两处复用:``GET /chat/interrupt``(刷新页面后补渲染卡片)与
  679. ``POST /chat/resume``(在构造 decisions 前先知道有几个待审批动作——
  680. 官方 HITL 要求 decisions 与 action_requests **数量严格一一对应**,
  681. 数量对不上会直接抛错)。
  682. """
  683. get_state = getattr(agent, "aget_state", None)
  684. if get_state is None:
  685. return []
  686. try:
  687. state = await get_state(config)
  688. except Exception as exc:
  689. logger.warning("读取图状态失败:%s: %s", type(exc).__name__, exc)
  690. raise
  691. found: list[Any] = []
  692. for task in getattr(state, "tasks", None) or ():
  693. for item in getattr(task, "interrupts", None) or ():
  694. found.append(pack(item))
  695. return found