136 lines
5.4 KiB
Python
136 lines
5.4 KiB
Python
"""异步执行与恢复大屏任务，当前使用单进程内存检查点"""
|
|
|
|
from contextlib import aclosing
|
|
from uuid import uuid4
|
|
|
|
from langgraph.types import Command
|
|
from loguru import logger
|
|
|
|
from app.core.exceptions import BizException
|
|
from app.observability.tracing import astream_screen
|
|
from app.schemas.agent import AgentResult, ResumeRequest
|
|
from app.schemas.request import DesignRequest
|
|
|
|
|
|
class AgentService:
|
|
def __init__(self, graph):
|
|
self.graph = graph
|
|
self._active: set[str] = set()
|
|
|
|
def _prepare_run(self, request: DesignRequest):
|
|
# 新请求始终独立建任务，不用 conversation_id 充当检查点 ID
|
|
thread_id = str(uuid4())
|
|
# 暂无会话创建接口，未传会话 ID 时先分配标识，不创建数据库记录
|
|
conversation_id = request.context.conversation_id or str(uuid4())
|
|
payload = request.model_dump(mode="json")
|
|
payload["context"]["conversation_id"] = conversation_id
|
|
|
|
return (
|
|
thread_id,
|
|
{
|
|
"request": payload,
|
|
"conversation_id": conversation_id,
|
|
"thread_id": thread_id,
|
|
"message_id": str(uuid4()),
|
|
},
|
|
conversation_id,
|
|
)
|
|
|
|
def stream_run(self, request: DesignRequest):
|
|
execution = self._prepare_run(request)
|
|
self._claim(execution[0])
|
|
return execution
|
|
|
|
async def prepare_resume(self, request: ResumeRequest):
|
|
# 在发送流响应头之前完成校验并占用任务
|
|
thread_id = str(request.thread_id)
|
|
# 检查和占用之间没有 await，同一事件循环内不会重复进入
|
|
self._claim(thread_id)
|
|
try:
|
|
snapshot = await self.graph.aget_state(
|
|
{"configurable": {"thread_id": thread_id}}
|
|
)
|
|
if not snapshot.values:
|
|
raise BizException(404, "任务不存在或已过期")
|
|
if not any(task.interrupts for task in snapshot.tasks):
|
|
raise BizException(409, "任务当前未等待用户补充")
|
|
# 恢复前校验选项，避免无效回答消耗一次 interrupt 恢复
|
|
clarification = snapshot.values["clarification"]
|
|
reply = request.reply
|
|
if (
|
|
not reply.cancel
|
|
and reply.option_id is not None
|
|
and reply.option_id not in {item.id for item in clarification.options}
|
|
):
|
|
raise BizException(422, "请选择当前问题的有效选项")
|
|
|
|
conversation_id = snapshot.values["conversation_id"]
|
|
# 恢复同一任务不能切换会话，补充编辑上下文时也保留原会话标识
|
|
if reply.context is not None:
|
|
if reply.context.conversation_id not in (None, conversation_id):
|
|
raise BizException(422, "恢复任务时不能切换会话")
|
|
reply = reply.model_copy(
|
|
update={
|
|
"context": reply.context.model_copy(
|
|
update={"conversation_id": conversation_id}
|
|
)
|
|
}
|
|
)
|
|
|
|
# 每次有效补充或取消生成新的输入 ID，任务和会话保持不变
|
|
return (
|
|
thread_id,
|
|
Command(
|
|
resume=reply.model_dump(mode="json"),
|
|
update={"message_id": str(uuid4())},
|
|
),
|
|
conversation_id,
|
|
)
|
|
except BaseException:
|
|
self.release(thread_id)
|
|
raise
|
|
|
|
def release(self, thread_id):
|
|
self._active.discard(thread_id)
|
|
|
|
def _claim(self, thread_id: str) -> None:
|
|
if thread_id in self._active:
|
|
raise BizException(409, "任务正在执行，请勿重复提交")
|
|
self._active.add(thread_id)
|
|
|
|
@staticmethod
|
|
def _result(state, thread_id):
|
|
waiting = bool(state.get("__interrupt__"))
|
|
|
|
# 只交付校验通过的完整动作批次，失败草稿留在后端
|
|
return AgentResult(
|
|
conversation_id=state["conversation_id"],
|
|
thread_id=thread_id,
|
|
message_id=state["message_id"],
|
|
status="waiting_for_input" if waiting else state["status"],
|
|
actions=state.get("actions") or []
|
|
if state.get("validation_passed") is True and state["status"] == "completed"
|
|
else [],
|
|
answer=state.get("answer"),
|
|
clarification=state.get("clarification") if waiting else None,
|
|
failure_reason=state.get("failure_reason"),
|
|
)
|
|
|
|
async def stream(self, thread_id, payload, session_id):
|
|
"""消费图事件，完整状态只用于整理最终结果，不直接暴露给客户端"""
|
|
state = None
|
|
with logger.contextualize(thread_id=thread_id):
|
|
async with aclosing(
|
|
astream_screen(
|
|
self.graph, payload, thread_id=thread_id, session_id=session_id
|
|
)
|
|
) as events:
|
|
async for mode, data in events:
|
|
if mode == "values":
|
|
state = data
|
|
elif data.get("type") == "node_progress" and data.get("label"):
|
|
yield "progress", data
|
|
if state is None:
|
|
raise RuntimeError("图未返回状态")
|
|
yield "final", self._result(state, thread_id).model_dump(mode="json")
|