jiawei/my-v0-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat：接入sse

Browse Source
This commit is contained in:
jiawei
2026-09-29 16:41:52 +08:00
parent 2b029b4bc7
commit e332259b07
4 changed files with 152 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/api/agent.py
+11 -8
View File
@@ -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)
app/common/sse.py
+65
View File
@@ -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])
app/observability/tracing.py
+38 -15
View File
@@ -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]:
app/services/agent.py
+38 -34
View File
@@ -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")