Files
my-v0-server/app/services/agent.py
T

228 lines
9.7 KiB
Python

"""异步执行与恢复大屏任务,当前使用单进程内存检查点"""
from contextlib import aclosing
from datetime import datetime, timezone
from uuid import uuid4
from langgraph.types import Command
from loguru import logger
from app.config import settings
from app.core.exceptions import BizException
from app.datafiles.service import DataFileService
from app.observability.tracing import astream_screen
from app.processor.history import trim_history
from app.schemas.agent import AgentResult, ResumeRequest
from app.schemas.conversation import Turn
from app.schemas.request import DesignRequest
from app.services.conversation import ConversationStore
class AgentService:
def __init__(self, graph, store: ConversationStore, data_files: DataFileService):
self.graph = graph
self.store = store
self.data_files = data_files
# 同一时刻,一个任务和一个会话都只允许一个执行,避免历史顺序错乱
self._active: set[str] = set()
self._conversation_of: dict[str, str] = {}
async def prepare_run(self, request: DesignRequest, user_id: int):
# 会话 ID 必须由后端签发,不再为缺失的 ID 兜底生成
conversation_id = request.context.conversation_id
if conversation_id is None:
raise BizException(422, "请先创建会话并传入 conversation_id")
if not await self.store.exists(conversation_id, user_id):
raise BizException(404, "会话不存在或已过期")
# 新请求始终独立建任务,不用 conversation_id 充当检查点 ID
thread_id = str(uuid4())
self._claim(thread_id, conversation_id)
try:
turns = await self.store.load_turns(
conversation_id, settings.conversation_history_turns
)
# 附件引用的文件不属于本会话时在流开始前返回 404
data_files = await self.data_files.request_files(
user_id, conversation_id, request.attachments
)
except BaseException:
self.release(thread_id)
raise
payload = request.model_dump(mode="json")
# 历史只信任服务端记录,覆盖前端可能携带的同名扩展字段;记录是完整的,这里按 token 预算裁剪
payload["history"] = trim_history(
[
{
"user": turn.prompt,
"assistant": turn.summary,
"outcome": turn.outcome,
}
for turn in turns
],
settings.history_max_tokens,
)
# 同样只信任服务端读取的结果;会话内的文件每一轮都可用,不附带也能查询
payload["data_files"] = [item.model_dump(mode="json") for item in data_files]
return (
thread_id,
{
"request": payload,
"conversation_id": conversation_id,
"thread_id": thread_id,
"message_id": str(uuid4()),
},
conversation_id,
user_id,
)
async def prepare_resume(self, request: ResumeRequest, user_id: int):
# 在发送流响应头之前完成校验并占用任务
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 or not await self.store.exists(
snapshot.values["conversation_id"], user_id
):
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"]
self._claim_conversation(thread_id, 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,
user_id,
)
except BaseException:
self.release(thread_id)
raise
def release(self, thread_id):
self._active.discard(thread_id)
self._conversation_of.pop(thread_id, None)
def _claim(self, thread_id: str, conversation_id: str | None = None) -> None:
if thread_id in self._active:
raise BizException(409, "任务正在执行,请勿重复提交")
self._active.add(thread_id)
if conversation_id is not None:
try:
self._claim_conversation(thread_id, conversation_id)
except BaseException:
self.release(thread_id)
raise
def _claim_conversation(self, thread_id: str, conversation_id: str) -> None:
if conversation_id in self._conversation_of.values():
raise BizException(409, "该会话有任务正在执行,请稍后再试")
self._conversation_of[thread_id] = conversation_id
@staticmethod
def _summarize(result: AgentResult) -> str:
"""本轮助手一侧的完整文本,只陈述后端生成了什么,不代表前端已执行;不截断,带入请求时再按预算裁剪"""
if result.status == "waiting_for_input":
question = result.clarification.question if result.clarification else ""
return f"向用户提问:{question}"
if result.status == "cancelled":
return "用户取消了本次补充"
if result.status == "failed":
return f"处理失败:{result.answer or '未知原因'}"
if result.answer:
return result.answer
return f"已生成 {len(result.actions)} 个画布动作"
async def _save_turn(self, state, result: AgentResult) -> None:
# 历史写入失败不影响本次结果交付,只记录日志
try:
await self.store.save_turn(
Turn(
conversation_id=result.conversation_id,
thread_id=result.thread_id,
message_id=result.message_id,
prompt=state["request"]["prompt"],
outcome=result.status,
summary=self._summarize(result),
created_at=datetime.now(timezone.utc),
)
)
except Exception:
logger.exception("保存对话轮次失败")
@staticmethod
def _result(state, thread_id):
waiting = bool(state.get("__interrupt__"))
# 只交付完成状态下的动作;actions 是 agent 最后一次成功提交、已通过模拟执行的整批动作
deliver = state["status"] == "completed"
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 deliver else [],
answer=state.get("answer"),
clarification=state.get("clarification") if waiting else None,
)
async def stream(self, thread_id, payload, session_id, user_id=None):
"""消费图事件,完整状态只用于整理最终结果,不直接暴露给客户端"""
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,
user_id=user_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
elif data.get("type") == "answer_delta":
# 问答文本逐段推送,完整回答仍在 final 事件中返回
yield "answer_delta", data
if state is None:
raise RuntimeError("图未返回状态")
result = self._result(state, thread_id)
await self._save_turn(state, result)
yield "final", result.model_dump(mode="json")