"""异步执行与恢复大屏任务，当前使用单进程内存检查点""" 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.base import JsonObject 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( [self._history_turn(turn) 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 @staticmethod def _history_turn(turn: Turn) -> JsonObject: """一轮历史：原话、回复、结果，有轮次记录时一并带上，由各节点决定给模型看哪些""" item: JsonObject = { "user": turn.prompt, "assistant": turn.summary, "outcome": turn.outcome, } for key in ("changes", "queries"): if records := getattr(turn, key): item[key] = [record.model_dump(mode="json") for record in records] if turn.clarifications: item["clarifications"] = turn.clarifications return item 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: # 历史写入失败不影响本次结果交付，只记录日志 completed = result.status == "completed" context = state.get("input_context") 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), # 只有完成的轮次交付了修改、问答成功；其余结果下状态里即使有也不记 changes=(state.get("changes") or []) if completed else [], queries=(state.get("queries") or []) if completed else [], clarifications=context.clarification_history if context else [], ) ) 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")