79 lines
2.0 KiB
Python
79 lines
2.0 KiB
Python
|
"""
|
|||
|
|
流式处理模块 - 处理 Agent 执行的流式输出
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
from typing import AsyncGenerator, Dict, Any
|
|||
|
|
|
|||
|
|
from backend.app.logger import info, error
|
|||
|
|
from .stream_context import set_stream_queue
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def run_graph_stream(
|
|||
|
|
graph,
|
|||
|
|
input_state: Dict[str, Any],
|
|||
|
|
config: Dict[str, Any],
|
|||
|
|
) -> AsyncGenerator[Dict[str, Any], None]:
|
|||
|
|
"""
|
|||
|
|
运行图并通过队列流式输出事件
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
graph: 编译后的 LangGraph
|
|||
|
|
input_state: 输入状态
|
|||
|
|
config: 配置
|
|||
|
|
|
|||
|
|
Yields:
|
|||
|
|
流式事件
|
|||
|
|
"""
|
|||
|
|
queue: asyncio.Queue = asyncio.Queue()
|
|||
|
|
set_stream_queue(queue)
|
|||
|
|
|
|||
|
|
async def run_graph():
|
|||
|
|
"""后台任务：运行 graph"""
|
|||
|
|
try:
|
|||
|
|
info(f"📡 开始调用 graph.astream()...")
|
|||
|
|
async for _ in graph.astream(
|
|||
|
|
input_state,
|
|||
|
|
config=config,
|
|||
|
|
stream_mode=["updates"],
|
|||
|
|
version="v2",
|
|||
|
|
subgraphs=True
|
|||
|
|
):
|
|||
|
|
# 流式事件都从 agent.py 节点内部通过队列发送了
|
|||
|
|
pass
|
|||
|
|
except Exception as e:
|
|||
|
|
error(f"❌ 执行图时出错: {e}")
|
|||
|
|
import traceback
|
|||
|
|
error(f"📋 堆栈: {traceback.format_exc()}")
|
|||
|
|
await queue.put({"type": "error", "message": str(e)})
|
|||
|
|
finally:
|
|||
|
|
await queue.put(None) # 结束哨兵
|
|||
|
|
|
|||
|
|
# 启动后台任务
|
|||
|
|
bg_task = asyncio.create_task(run_graph())
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
while True:
|
|||
|
|
event = await queue.get()
|
|||
|
|
if event is None:
|
|||
|
|
break
|
|||
|
|
yield event
|
|||
|
|
|
|||
|
|
except GeneratorExit:
|
|||
|
|
info("⚠️ GeneratorExit，取消后台任务")
|
|||
|
|
bg_task.cancel()
|
|||
|
|
raise
|
|||
|
|
finally:
|
|||
|
|
await _cleanup_task(bg_task)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def _cleanup_task(bg_task: asyncio.Task) -> None:
|
|||
|
|
"""清理后台任务"""
|
|||
|
|
if not bg_task.done():
|
|||
|
|
info("⏹️ 清理后台任务")
|
|||
|
|
bg_task.cancel()
|
|||
|
|
try:
|
|||
|
|
await bg_task
|
|||
|
|
except asyncio.CancelledError:
|
|||
|
|
info("✅ 后台任务已取消")
|