113 lines
4.3 KiB
Python
113 lines
4.3 KiB
Python
from collections.abc import AsyncIterator
|
|
from dataclasses import asdict
|
|
|
|
from fastapi import APIRouter, Depends, Request, status
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from starlette.responses import StreamingResponse
|
|
|
|
from app.api.dependencies.auth import build_request_context
|
|
from app.api.dependencies.database import get_session
|
|
from app.api.dependencies.rate_limit import enforce_rate_limit
|
|
from app.api.schemas.agent_runs import (
|
|
AgentRunAcceptedEnvelope,
|
|
AgentRunAcceptedResponse,
|
|
AgentRunCreateRequest,
|
|
AgentRunStatusEnvelope,
|
|
AgentRunStatusResponse,
|
|
)
|
|
from app.api.views.agent_run_sse import encode_events, recovery_events
|
|
from app.core.config import get_settings
|
|
from app.core.contracts import AgentRequest, RequestContext
|
|
from app.core.errors import SseNotAcceptableError
|
|
from app.service.agent_run_application_service import AgentRunApplicationService
|
|
from app.service.run_query_service import RunQueryService
|
|
|
|
router = APIRouter(
|
|
prefix="/api/v1/agent-runs",
|
|
tags=["agent-runs"],
|
|
dependencies=[Depends(enforce_rate_limit)],
|
|
)
|
|
|
|
|
|
def accepts_event_stream(accept: str | None) -> bool:
|
|
if accept is None or not accept.strip():
|
|
return True
|
|
for entry in accept.split(","):
|
|
parts = entry.split(";")
|
|
if parts[0].strip().lower() not in {"text/event-stream", "text/*", "*/*"}:
|
|
continue
|
|
quality = 1.0
|
|
for parameter in parts[1:]:
|
|
name, _, value = parameter.partition("=")
|
|
if name.strip().lower() == "q":
|
|
try:
|
|
quality = float(value.strip())
|
|
except ValueError:
|
|
quality = 0.0
|
|
if quality > 0:
|
|
return True
|
|
return False
|
|
|
|
|
|
@router.post("", response_model=AgentRunAcceptedEnvelope, status_code=status.HTTP_202_ACCEPTED)
|
|
async def create_agent_run(
|
|
payload: AgentRunCreateRequest,
|
|
request: Request,
|
|
context: RequestContext = Depends(build_request_context), # noqa: B008
|
|
session: AsyncSession = Depends(get_session), # noqa: B008
|
|
) -> AgentRunAcceptedEnvelope:
|
|
request.state.request_context = context
|
|
accepted = await AgentRunApplicationService(session).accept(
|
|
AgentRequest(**payload.model_dump()), context)
|
|
return AgentRunAcceptedEnvelope(
|
|
data=AgentRunAcceptedResponse(
|
|
run_id=accepted.run_id, trace_id=accepted.trace_id, status=accepted.status,
|
|
status_url=f"/api/v1/agent-runs/{accepted.run_id}",
|
|
events_url=f"/api/v1/agent-runs/{accepted.run_id}/events",
|
|
),
|
|
meta={"trace_id": context.trace_id},
|
|
)
|
|
|
|
|
|
@router.get("/{run_id}", response_model=AgentRunStatusEnvelope)
|
|
async def get_agent_run(
|
|
run_id: str, context: RequestContext = Depends(build_request_context), # noqa: B008
|
|
) -> AgentRunStatusEnvelope:
|
|
return AgentRunStatusEnvelope(
|
|
data=AgentRunStatusResponse(**asdict(await RunQueryService().get(run_id, context))),
|
|
meta={"trace_id": context.trace_id},
|
|
)
|
|
|
|
|
|
@router.get("/{run_id}/events")
|
|
async def stream_agent_run_events(
|
|
run_id: str,
|
|
request: Request,
|
|
context: RequestContext = Depends(build_request_context), # noqa: B008
|
|
) -> StreamingResponse:
|
|
query = RunQueryService()
|
|
initial = await query.get(run_id, context)
|
|
if not accepts_event_stream(request.headers.get("Accept")):
|
|
raise SseNotAcceptableError("Accept must allow text/event-stream")
|
|
|
|
async def generate() -> AsyncIterator[str]:
|
|
start_sent = False
|
|
async for snapshot in query.watch(initial, context):
|
|
if snapshot is None:
|
|
yield ": heartbeat\n\n"
|
|
continue
|
|
result = snapshot.result or {}
|
|
events = recovery_events(
|
|
run_id=snapshot.run_id, trace_id=snapshot.trace_id, status=snapshot.status,
|
|
error_code=snapshot.error_code, content=result.get("content"),
|
|
tool_calls=result.get("tool_calls"),
|
|
replay=initial.status in {"succeeded", "failed", "cancelled"},
|
|
chunk_size=get_settings().sse_chunk_characters,
|
|
)
|
|
for encoded in encode_events(run_id, events[1:] if start_sent else events):
|
|
yield encoded
|
|
start_sent = True
|
|
|
|
return StreamingResponse(generate(), media_type="text/event-stream",
|
|
headers={"Cache-Control": "no-cache"})
|