Files

75 lines
3.1 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.schemas.agent_runs import (
AgentRunAcceptedResponse,
AgentRunCreateRequest,
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.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"])
@router.post("", response_model=AgentRunAcceptedResponse, 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
) -> AgentRunAcceptedResponse:
request.state.request_context = context
accepted = await AgentRunApplicationService(session).accept(
AgentRequest(**payload.model_dump()), context)
return 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",
)
@router.get("/{run_id}", response_model=AgentRunStatusResponse)
async def get_agent_run(
run_id: str, context: RequestContext = Depends(build_request_context), # noqa: B008
) -> AgentRunStatusResponse:
return AgentRunStatusResponse(**asdict(await RunQueryService().get(run_id, context)))
@router.get("/{run_id}/events")
async def stream_agent_run_events(
run_id: str, context: RequestContext = Depends(build_request_context), # noqa: B008
) -> StreamingResponse:
query = RunQueryService()
initial = await query.get(run_id, context)
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"})