91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
"""通过 LangGraph custom 流发送进度，不额外保存任务状态"""
|
|
|
|
import asyncio
|
|
from contextlib import contextmanager
|
|
from time import perf_counter
|
|
from uuid import uuid4
|
|
|
|
from langgraph.config import get_config, get_stream_writer
|
|
from langgraph.errors import GraphInterrupt
|
|
|
|
from app.schemas.node_progress import NodeProgress, NodeProgressStatus
|
|
|
|
# 将内部节点名映射为用户能理解的进度说明，未配置的节点仍发送事件但不附带 label
|
|
NODE_PROGRESS_MESSAGES = {
|
|
"requirement_node": "正在理解你的需求",
|
|
"qa_node": "正在整理回答",
|
|
"component_resource_node": "正在选择合适的图表和组件",
|
|
"planner_node": "正在设计布局",
|
|
"action_node": "正在准备画布修改",
|
|
"validator_node": "正在检查修改结果",
|
|
"reflection_node": "正在调整修改方案",
|
|
}
|
|
|
|
|
|
@contextmanager
|
|
def node_progress(node: str, state):
|
|
"""每次调用独立计时和编号，同步与异步入口共用"""
|
|
|
|
try:
|
|
# 读取当前图执行配置，用于判断是否处于图上下文，并获取任务标识
|
|
config = get_config()
|
|
except RuntimeError:
|
|
# 独立调用节点没有图上下文，直接执行业务逻辑
|
|
# 将控制权交给 with 块中的节点逻辑，此分支不发送进度事件
|
|
yield
|
|
# 节点逻辑正常执行后直接退出，不再获取图的流写入器
|
|
return
|
|
|
|
# 获取当前图的 custom 流写入器，供本次节点执行推送进度事件
|
|
writer = get_stream_writer()
|
|
# 以实际检查点配置为准，兼容没有在 state 中填写 ID 的本地示例
|
|
thread_id = config.get("configurable", {}).get("thread_id") or state.get(
|
|
"thread_id"
|
|
)
|
|
# 为本次节点调用生成唯一编号，区分同一节点的首次执行、重试和恢复执行
|
|
execution_id = str(uuid4())
|
|
# 使用单调计时器记录起点，后续计算本次执行耗时
|
|
started = perf_counter()
|
|
|
|
def emit(status: NodeProgressStatus):
|
|
# 将任务、用户消息和本次执行编号放入同一事件，便于消费方关联进度
|
|
event = NodeProgress(
|
|
thread_id=thread_id,
|
|
message_id=state.get("message_id"),
|
|
execution_id=execution_id,
|
|
node=node,
|
|
# 未配置用户说明时返回 None，序列化时省略该字段
|
|
label=NODE_PROGRESS_MESSAGES.get(node),
|
|
status=status,
|
|
# 开始事件不携带耗时，其他状态携带从本次调用开始计算的毫秒数
|
|
duration_ms=None
|
|
if status == "started"
|
|
else round((perf_counter() - started) * 1000, 3),
|
|
)
|
|
|
|
# 转成可序列化的字典并发送到 custom 流，省略空字段，不写入业务状态
|
|
writer(event.model_dump(mode="json", exclude_none=True))
|
|
|
|
# 节点业务逻辑开始前，先通知消费方本次执行已经启动
|
|
emit("started")
|
|
try:
|
|
# 执行 with 块中的节点逻辑，节点抛出的异常会回到下面对应的异常分支
|
|
yield
|
|
except GraphInterrupt:
|
|
# interrupt 是正常暂停，恢复后会产生新的执行 ID
|
|
emit("waiting_for_input")
|
|
# 将暂停信号重新抛给 LangGraph，由图负责保存检查点和中断执行
|
|
raise
|
|
except asyncio.CancelledError:
|
|
# 协程被取消也要结束这次节点执行，继续向上传递取消信号
|
|
emit("cancelled")
|
|
raise
|
|
except Exception:
|
|
# 只通知失败，异常详情仍由原有日志和 API 异常处理负责
|
|
emit("failed")
|
|
raise
|
|
else:
|
|
# 校验不通过等业务结果仍属于节点正常执行完成
|
|
# completed 仅表示本次节点调用未抛异常，不代表整个任务成功或前端已执行
|
|
emit("completed")
|