diff --git a/app/api/agent.py b/app/api/agent.py index 6ac3d9e..06a01d9 100644 --- a/app/api/agent.py +++ b/app/api/agent.py @@ -2,18 +2,21 @@ from fastapi import APIRouter, Request -from app.common.response import ApiResponse, success_response -from app.schemas.agent import AgentResult, ResumeRequest +from app.common.sse import AgentEventResponse +from app.schemas.agent import ResumeRequest from app.schemas.request import DesignRequest router = APIRouter(prefix="/agent", tags=["Agent"]) -@router.post("/run", response_model=ApiResponse[AgentResult]) -async def run_agent(body: DesignRequest, request: Request): - return success_response(await request.app.state.agent_service.run(body)) +@router.post("/run", response_class=AgentEventResponse, status_code=200) +async def stream_agent(body: DesignRequest, request: Request): + service = request.app.state.agent_service + return AgentEventResponse(service, service.stream_run(body), request) -@router.post("/resume", response_model=ApiResponse[AgentResult]) -async def resume_agent(body: ResumeRequest, request: Request): - return success_response(await request.app.state.agent_service.resume(body)) +@router.post("/resume", response_class=AgentEventResponse, status_code=200) +async def stream_resume_agent(body: ResumeRequest, request: Request): + service = request.app.state.agent_service + execution = await service.prepare_resume(body) + return AgentEventResponse(service, execution, request) diff --git a/app/common/sse.py b/app/common/sse.py new file mode 100644 index 0000000..3f1cff8 --- /dev/null +++ b/app/common/sse.py @@ -0,0 +1,65 @@ +"""SSE 编码和连接生命周期，不保存全局任务队列""" + +import json +from contextlib import aclosing + +import anyio +from starlette.responses import StreamingResponse + +from app.common.response import error_content, success_response +from app.core.exceptions import BizException, handle_unexpected_exception + + +def pack_event(event, content): + return f"event: {event}\ndata: {json.dumps(content, ensure_ascii=False)}\n\n" + + +async def encode_events(service, execution, request): + try: + async with aclosing(service.stream(*execution)) as events: + async for event, data in events: + yield pack_event(event, success_response(data).model_dump(mode="json")) + except BizException as exc: + yield pack_event("error", error_content(exc.code, exc.msg)) + except Exception as exc: # noqa: BLE001 流中异常统一封装后终止 + # 流开始后不能修改 HTTP 状态，复用统一异常映射并发送终止事件 + exc._agent_thread_id = execution[0] + response = await handle_unexpected_exception(request, exc) + yield pack_event("error", json.loads(response.body)) + + +class AgentEventResponse(StreamingResponse): + media_type = "text/event-stream" + + def __init__(self, service, execution, request): + self.service = service + self.execution = execution + super().__init__( + encode_events(service, execution, request), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) + + async def stream_response(self, send): + try: + await super().stream_response(send) + finally: + # 在消费流的同一协程关闭生成器，避免追踪上下文跨协程退出 + with anyio.CancelScope(shield=True): + await self.body_iterator.aclose() + + async def __call__(self, scope, receive, send): + try: + # 同时监听断开连接，模型尚未产生下一个事件时也能及时取消 + async with anyio.create_task_group() as group: + + async def stream(): + await self.stream_response(send) + group.cancel_scope.cancel() + + group.start_soon(stream) + await self.listen_for_disconnect(receive) + group.cancel_scope.cancel() + finally: + # 即使连接在生成器开始前断开，也要释放预先占用的任务 + self.service.release(self.execution[0]) diff --git a/app/observability/tracing.py b/app/observability/tracing.py index a34649b..f16eef5 100644 --- a/app/observability/tracing.py +++ b/app/observability/tracing.py @@ -67,10 +67,43 @@ async def ainvoke_screen( session_id 由调用方在首次与恢复时一致传入（通常为 conversation_id）。 未传时始终使用 thread_id，不从仅首次请求才有的 payload 推断。 """ + with _screen_trace(payload, thread_id, session_id) as (config, root): + result = await graph.ainvoke(payload, config=config) + _finish_trace(root, result) + return result + + +async def astream_screen(graph, payload, *, thread_id, session_id=None): + """同一轮流保持同一个追踪上下文，关闭流时结束追踪""" + from contextlib import aclosing + + result = {} + with _screen_trace(payload, thread_id, session_id) as (config, root): + async with aclosing( + graph.astream(payload, config=config, stream_mode=["custom", "values"]) + ) as events: + async for mode, data in events: + if mode == "values": + result = data + yield mode, data + _finish_trace(root, result) + + +def _finish_trace(root, result): + if root is not None: + output = _run_output(result) + root.update( + output=output, level="ERROR" if output["outcome"] == "failed" else "DEFAULT" + ) + + +@contextmanager +def _screen_trace(payload, thread_id, session_id): config = {"configurable": {"thread_id": thread_id}} client = _ensure_client() if tracing_enabled() else None if client is None: - return await graph.ainvoke(payload, config=config) + yield config, None + return from langfuse import propagate_attributes from langfuse.langchain import CallbackHandler @@ -103,15 +136,11 @@ async def ainvoke_screen( ), ): try: - result = await graph.ainvoke( - payload, - config={ - **config, - "callbacks": [handler], - "run_name": "run-screen-graph", - }, + yield ( + {**config, "callbacks": [handler], "run_name": "run-screen-graph"}, + root, ) - except CancelledError: + except (CancelledError, GeneratorExit): root.update(output={"outcome": "cancelled"}) raise except Exception as error: @@ -121,12 +150,6 @@ async def ainvoke_screen( status_message=str(error), ) raise - output = _run_output(result) - root.update( - output=output, - level="ERROR" if output["outcome"] == "failed" else "DEFAULT", - ) - return result def _run_output(state: dict[str, Any]) -> dict[str, Any]: diff --git a/app/services/agent.py b/app/services/agent.py index 222ae29..18bb9e6 100644 --- a/app/services/agent.py +++ b/app/services/agent.py @@ -1,12 +1,13 @@ """异步执行与恢复大屏任务，当前使用单进程内存检查点""" +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 ainvoke_screen +from app.observability.tracing import astream_screen from app.schemas.agent import AgentResult, ResumeRequest from app.schemas.request import DesignRequest @@ -16,7 +17,7 @@ class AgentService: self.graph = graph self._active: set[str] = set() - async def run(self, request: DesignRequest) -> AgentResult: + def _prepare_run(self, request: DesignRequest): # 新请求始终独立建任务，不用 conversation_id 充当检查点 ID thread_id = str(uuid4()) # 暂无会话创建接口，未传会话 ID 时先分配标识，不创建数据库记录 @@ -24,7 +25,7 @@ class AgentService: payload = request.model_dump(mode="json") payload["context"]["conversation_id"] = conversation_id - return await self._invoke( + return ( thread_id, { "request": payload, @@ -35,7 +36,13 @@ class AgentService: conversation_id, ) - async def resume(self, request: ResumeRequest) -> AgentResult: + 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) @@ -71,7 +78,7 @@ class AgentService: ) # 每次有效补充或取消生成新的输入 ID，任务和会话保持不变 - return await self._execute( + return ( thread_id, Command( resume=reply.model_dump(mode="json"), @@ -79,41 +86,20 @@ class AgentService: ), conversation_id, ) - finally: - self._active.discard(thread_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) - 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"), - ) - + @staticmethod + def _result(state, thread_id): waiting = bool(state.get("__interrupt__")) # 只交付校验通过的完整动作批次，失败草稿留在后端 @@ -129,3 +115,21 @@ class AgentService: 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")