254 lines
11 KiB
Python
254 lines
11 KiB
Python
"""异步执行与恢复大屏任务，当前使用单进程内存检查点"""
|
|
|
|
from contextlib import aclosing
|
|
from datetime import UTC, datetime
|
|
from uuid import uuid4
|
|
|
|
from langgraph.types import Command
|
|
from loguru import logger
|
|
|
|
from app.agent.data_tools import DataContext
|
|
from app.config import settings
|
|
from app.core.exceptions import BizException
|
|
from app.datafiles.insights import ready_insights
|
|
from app.datafiles.service import DataFileService
|
|
from app.observability.tracing import astream_screen
|
|
from app.processor.history import trim_history
|
|
from app.processor.state import RunContext
|
|
from app.schemas.agent import AgentResult, ResumeRequest
|
|
from app.schemas.conversation import Turn
|
|
from app.schemas.data_file import DataFileRecord, RequestDataFile
|
|
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
|
|
attached = request.data_file_ids
|
|
records = await self.data_files.request_files(
|
|
user_id, conversation_id, attached
|
|
)
|
|
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"] = [
|
|
RequestDataFile(
|
|
file_id=record.file_id,
|
|
name=record.name,
|
|
attached_this_turn=record.file_id in attached,
|
|
tables=record.tables,
|
|
insights=ready_insights(record.insights),
|
|
).model_dump(mode="json")
|
|
for record in records
|
|
]
|
|
|
|
return (
|
|
thread_id,
|
|
{
|
|
"request": payload,
|
|
"conversation_id": conversation_id,
|
|
"thread_id": thread_id,
|
|
"message_id": str(uuid4()),
|
|
},
|
|
conversation_id,
|
|
user_id,
|
|
self._run_context(records),
|
|
)
|
|
|
|
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}
|
|
)
|
|
}
|
|
)
|
|
# 运行上下文不在检查点里，恢复时按会话重新读取；期间删除或过期的文件工具查不到
|
|
records = await self.data_files.request_files(
|
|
user_id, conversation_id, set()
|
|
)
|
|
|
|
# 每次有效补充或取消生成新的输入 ID，任务和会话保持不变
|
|
return (
|
|
thread_id,
|
|
Command(
|
|
resume=reply.model_dump(mode="json"),
|
|
update={"message_id": str(uuid4())},
|
|
),
|
|
conversation_id,
|
|
user_id,
|
|
self._run_context(records),
|
|
)
|
|
except BaseException:
|
|
self.release(thread_id)
|
|
raise
|
|
|
|
def _run_context(self, records: list[DataFileRecord]) -> RunContext:
|
|
return RunContext(data=DataContext(tuple(records), self.data_files.root))
|
|
|
|
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(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,
|
|
suggestions=(state.get("suggestions") or []) if deliver else [],
|
|
)
|
|
|
|
async def stream(self, thread_id, payload, session_id, user_id=None, context=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,
|
|
context=context,
|
|
)
|
|
) 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")
|