feat: AG-UI 协议对齐与路由导航功能
- 前端: 添加 SSE 流式支持、stateSnapshot 事件、路由导航工具 - 前端: 实现工具调用审批流程，支持 pending 状态展示 - 后端: Agent 状态管理与会话持久化相关重构 - 文档: 新增 agent-agui-full-alignance 设计文档 - 测试: 补充相关单元测试和集成测试
This commit is contained in:
zl-q
+141
-32
@@ -2,96 +2,177 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
import asyncio
|
||||
from datetime import date
|
||||
import re
|
||||
import time
|
||||
from typing import Annotated
|
||||
|
||||
from ag_ui.core import RunAgentInput
|
||||
from fastapi import APIRouter, Depends, Header, Query, Request, status
|
||||
from fastapi import HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from core.agent.infrastructure.agui.stream import to_sse_event
|
||||
from core.agent.domain.agui_input import parse_run_input
|
||||
from core.auth.models import CurrentUser
|
||||
from services.base.redis import get_or_init_redis_client
|
||||
from v1.agent.dependencies import get_agent_service
|
||||
from v1.agent.schemas import ResumeRequest, RunRequest, TaskAcceptedResponse
|
||||
from v1.agent.schemas import TaskAcceptedResponse
|
||||
from v1.agent.service import AgentService
|
||||
from v1.users.dependencies import get_current_user
|
||||
|
||||
router = APIRouter(prefix="/agent", tags=["agent"])
|
||||
_LAST_EVENT_ID_RE = re.compile(r"^\d+-\d+$")
|
||||
_RUNS_PER_MINUTE = 30
|
||||
_MAX_SSE_CONNECTIONS_PER_USER = 3
|
||||
_SSE_SLOT_TTL_SECONDS = 15 * 60
|
||||
|
||||
|
||||
async def _allow_run_request(*, user_id: str) -> bool:
|
||||
try:
|
||||
redis = await get_or_init_redis_client()
|
||||
minute_bucket = int(time.time() // 60)
|
||||
key = f"agent:run-rate:{user_id}:{minute_bucket}"
|
||||
count = await redis.incr(key)
|
||||
if count == 1:
|
||||
await redis.expire(key, 70)
|
||||
return int(count) <= _RUNS_PER_MINUTE
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
|
||||
async def _acquire_sse_slot(*, user_id: str) -> bool:
|
||||
try:
|
||||
redis = await get_or_init_redis_client()
|
||||
key = f"agent:sse-active:{user_id}"
|
||||
count = await redis.incr(key)
|
||||
if count == 1:
|
||||
await redis.expire(key, _SSE_SLOT_TTL_SECONDS)
|
||||
if int(count) > _MAX_SSE_CONNECTIONS_PER_USER:
|
||||
await redis.decr(key)
|
||||
return False
|
||||
return True
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
|
||||
async def _release_sse_slot(*, user_id: str) -> None:
|
||||
try:
|
||||
redis = await get_or_init_redis_client()
|
||||
key = f"agent:sse-active:{user_id}"
|
||||
count = await redis.decr(key)
|
||||
if int(count) <= 0:
|
||||
await redis.delete(key)
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
|
||||
|
||||
@router.post(
|
||||
"/runs", response_model=TaskAcceptedResponse, status_code=status.HTTP_202_ACCEPTED
|
||||
)
|
||||
async def enqueue_run(
|
||||
request: RunRequest,
|
||||
request: RunAgentInput,
|
||||
service: Annotated[AgentService, Depends(get_agent_service)],
|
||||
current_user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
) -> TaskAcceptedResponse:
|
||||
try:
|
||||
parse_run_input(request.model_dump(mode="json", by_alias=True))
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
allowed = await _allow_run_request(user_id=str(current_user.id))
|
||||
if not allowed:
|
||||
raise HTTPException(status_code=429, detail="Too many run requests")
|
||||
|
||||
task = await service.enqueue_run(
|
||||
session_id=request.session_id,
|
||||
prompt=request.prompt,
|
||||
run_input=request,
|
||||
current_user=current_user,
|
||||
)
|
||||
return TaskAcceptedResponse(
|
||||
task_id=task.task_id,
|
||||
session_id=task.session_id,
|
||||
thread_id=task.thread_id,
|
||||
run_id=task.run_id,
|
||||
created=task.created,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/runs/{session_id}/resume",
|
||||
"/runs/{thread_id}/resume",
|
||||
response_model=TaskAcceptedResponse,
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def enqueue_resume(
|
||||
session_id: str,
|
||||
request: ResumeRequest,
|
||||
thread_id: str,
|
||||
request: RunAgentInput,
|
||||
service: Annotated[AgentService, Depends(get_agent_service)],
|
||||
current_user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
) -> TaskAcceptedResponse:
|
||||
if request.thread_id != thread_id:
|
||||
raise HTTPException(status_code=422, detail="thread_id path/body mismatch")
|
||||
try:
|
||||
parse_run_input(request.model_dump(mode="json", by_alias=True))
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
task = await service.enqueue_resume(
|
||||
session_id=session_id,
|
||||
tool_call_id=request.tool_call_id,
|
||||
thread_id=thread_id,
|
||||
run_input=request,
|
||||
current_user=current_user,
|
||||
)
|
||||
return TaskAcceptedResponse(
|
||||
task_id=task.task_id,
|
||||
session_id=task.session_id,
|
||||
thread_id=task.thread_id,
|
||||
run_id=task.run_id,
|
||||
created=task.created,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/runs/{session_id}/events")
|
||||
@router.get("/runs/{thread_id}/events")
|
||||
async def stream_events(
|
||||
request: Request,
|
||||
session_id: str,
|
||||
thread_id: str,
|
||||
service: Annotated[AgentService, Depends(get_agent_service)],
|
||||
current_user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
last_event_id: str | None = Header(default=None, alias="Last-Event-ID"),
|
||||
idle_limit: int = Query(default=300, ge=1, le=3600),
|
||||
) -> StreamingResponse:
|
||||
if (
|
||||
last_event_id is not None
|
||||
and (
|
||||
len(last_event_id) > 32
|
||||
or _LAST_EVENT_ID_RE.fullmatch(last_event_id) is None
|
||||
)
|
||||
):
|
||||
raise HTTPException(status_code=422, detail="Invalid Last-Event-ID")
|
||||
|
||||
sse_slot_acquired = await _acquire_sse_slot(user_id=str(current_user.id))
|
||||
if not sse_slot_acquired:
|
||||
raise HTTPException(status_code=429, detail="Too many SSE connections")
|
||||
|
||||
async def _event_iter() -> AsyncIterator[str]:
|
||||
cursor = last_event_id
|
||||
idle_polls = 0
|
||||
while not await request.is_disconnected() and idle_polls < idle_limit:
|
||||
rows = await service.stream_events(
|
||||
session_id=session_id,
|
||||
last_event_id=cursor,
|
||||
current_user=current_user,
|
||||
)
|
||||
if not rows:
|
||||
idle_polls += 1
|
||||
yield ": keep-alive\n\n"
|
||||
await asyncio.sleep(0.2)
|
||||
continue
|
||||
|
||||
idle_polls = 0
|
||||
for row in rows:
|
||||
row_id = str(row.get("id", ""))
|
||||
event = row.get("event")
|
||||
if not row_id or not isinstance(event, dict):
|
||||
try:
|
||||
while not await request.is_disconnected() and idle_polls < idle_limit:
|
||||
rows = await service.stream_events(
|
||||
thread_id=thread_id,
|
||||
last_event_id=cursor,
|
||||
current_user=current_user,
|
||||
)
|
||||
if not rows:
|
||||
idle_polls += 1
|
||||
yield ": keep-alive\n\n"
|
||||
await asyncio.sleep(0.2)
|
||||
continue
|
||||
cursor = row_id
|
||||
yield to_sse_event(row_id, event)
|
||||
|
||||
idle_polls = 0
|
||||
for row in rows:
|
||||
row_id = str(row.get("id", ""))
|
||||
event = row.get("event")
|
||||
if not row_id or not isinstance(event, dict):
|
||||
continue
|
||||
cursor = row_id
|
||||
yield to_sse_event(row_id, event)
|
||||
finally:
|
||||
await _release_sse_slot(user_id=str(current_user.id))
|
||||
|
||||
return StreamingResponse(
|
||||
_event_iter(),
|
||||
@@ -102,3 +183,31 @@ async def stream_events(
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/runs/{thread_id}/history")
|
||||
async def get_history_snapshot(
|
||||
thread_id: str,
|
||||
service: Annotated[AgentService, Depends(get_agent_service)],
|
||||
current_user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
before: date | None = Query(default=None),
|
||||
) -> dict[str, object]:
|
||||
return await service.get_history_snapshot(
|
||||
thread_id=thread_id,
|
||||
before=before,
|
||||
current_user=current_user,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/history")
|
||||
async def get_user_history_snapshot(
|
||||
service: Annotated[AgentService, Depends(get_agent_service)],
|
||||
current_user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
thread_id: str | None = Query(default=None, alias="threadId"),
|
||||
before: date | None = Query(default=None),
|
||||
) -> dict[str, object]:
|
||||
return await service.get_user_history_snapshot(
|
||||
current_user=current_user,
|
||||
thread_id=thread_id,
|
||||
before=before,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user