66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
"""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])
|