132 lines
5.4 KiB
Python
132 lines
5.4 KiB
Python
"""异步执行与恢复大屏任务，当前使用单进程内存检查点"""
|
|
|
|
from uuid import uuid4
|
|
|
|
from langgraph.types import Command
|
|
from loguru import logger
|
|
|
|
from app.core.exceptions import BizException
|
|
from app.observability.tracing import ainvoke_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()
|
|
|
|
async def run(self, request: DesignRequest) -> AgentResult:
|
|
# 新请求始终独立建任务，不用 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 await self._invoke(
|
|
thread_id,
|
|
{
|
|
"request": payload,
|
|
"conversation_id": conversation_id,
|
|
"thread_id": thread_id,
|
|
"message_id": str(uuid4()),
|
|
},
|
|
conversation_id,
|
|
)
|
|
|
|
async def resume(self, request: ResumeRequest) -> AgentResult:
|
|
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 await self._execute(
|
|
thread_id,
|
|
Command(
|
|
resume=reply.model_dump(mode="json"),
|
|
update={"message_id": str(uuid4())},
|
|
),
|
|
conversation_id,
|
|
)
|
|
finally:
|
|
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)
|
|
|
|
async def _invoke(self, thread_id, payload, session_id) -> AgentResult:
|
|
self._claim(thread_id)
|
|
try:
|
|
return await self._execute(thread_id, payload, session_id)
|
|
finally:
|
|
# 异常和取消也必须释放占用，不阻塞后续请求
|
|
self._active.discard(thread_id)
|
|
|
|
async def _execute(self, thread_id, payload, session_id) -> AgentResult:
|
|
# 上下文只作用于当前异步任务，节点自动继承任务标识
|
|
with logger.contextualize(thread_id=thread_id):
|
|
logger.info("任务执行开始")
|
|
try:
|
|
state = await ainvoke_screen(
|
|
self.graph, payload, thread_id=thread_id, session_id=session_id
|
|
)
|
|
except Exception as exc:
|
|
# API 异常处理发生在上下文退出后，保留任务标识用于定位
|
|
exc._agent_thread_id = thread_id
|
|
raise
|
|
logger.info(
|
|
"任务执行结束，状态 {}",
|
|
"waiting_for_input"
|
|
if state.get("__interrupt__")
|
|
else state.get("status"),
|
|
)
|
|
|
|
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"),
|
|
)
|