76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
"""通过 LangGraph custom 流发送进度，不额外保存任务状态"""
|
|
|
|
import asyncio
|
|
from collections.abc import Mapping
|
|
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
|
|
|
|
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:
|
|
# 独立调用节点没有图上下文，直接执行业务逻辑
|
|
yield
|
|
return
|
|
|
|
writer = get_stream_writer()
|
|
values = state if isinstance(state, Mapping) else {}
|
|
# 以实际检查点配置为准，兼容没有在 state 中填写 ID 的本地示例
|
|
thread_id = config.get("configurable", {}).get("thread_id") or values.get(
|
|
"thread_id"
|
|
)
|
|
execution_id = str(uuid4())
|
|
started = perf_counter()
|
|
|
|
def emit(status: NodeProgressStatus):
|
|
event = NodeProgress(
|
|
thread_id=thread_id,
|
|
message_id=values.get("message_id"),
|
|
execution_id=execution_id,
|
|
node=node,
|
|
label=NODE_PROGRESS_MESSAGES.get(node),
|
|
status=status,
|
|
duration_ms=None
|
|
if status == "started"
|
|
else round((perf_counter() - started) * 1000, 3),
|
|
)
|
|
writer(event.model_dump(mode="json", exclude_none=True))
|
|
|
|
emit("started")
|
|
try:
|
|
yield
|
|
except GraphInterrupt:
|
|
# interrupt 是正常暂停，恢复后会产生新的执行 ID
|
|
emit("waiting_for_input")
|
|
raise
|
|
except asyncio.CancelledError:
|
|
# 协程被取消也要结束这次节点执行，继续向上传递取消信号
|
|
emit("cancelled")
|
|
raise
|
|
except Exception:
|
|
# 只通知失败，异常详情仍由原有日志和 API 异常处理负责
|
|
emit("failed")
|
|
raise
|
|
else:
|
|
# 校验不通过等业务结果仍属于节点正常执行完成
|
|
emit("completed")
|