app/ai/api.py

from __future__ import annotations

from typing import Annotated

from fastapi import APIRouter, Header, Path, Query, Request, status

from app.ai.schemas import AiCapabilities, AiGenerationRequest, AiHistory, AiJob
from app.security.errors import ForbiddenError

router = APIRouter(prefix="/v1/ai", tags=["ai"])


def _identity(request: Request) -> tuple[str, str, str | None, str]:
    context = request.state.auth
    if not context.workspace_id or not context.user_id:
        raise ForbiddenError
    return (
        context.workspace_id,
        context.user_id,
        context.api_key_id,
        request.state.request_id,
    )


@router.get("/capabilities", response_model=AiCapabilities)
async def capabilities(request: Request) -> AiCapabilities:
    return request.app.state.container.ai.capabilities()


@router.get("/jobs", response_model=AiHistory)
async def list_jobs(
    request: Request,
    offset: Annotated[int, Query(ge=0)] = 0,
    limit: Annotated[int, Query(ge=1, le=100)] = 25,
) -> AiHistory:
    workspace_id, user_id, _, _ = _identity(request)
    return await request.app.state.container.ai.history(
        workspace_id=workspace_id,
        user_id=user_id,
        offset=offset,
        limit=limit,
    )


@router.post("/jobs", response_model=AiJob, status_code=status.HTTP_202_ACCEPTED)
async def create_job(
    request: Request,
    payload: AiGenerationRequest,
    idempotency_key: Annotated[str, Header(alias="Idempotency-Key", min_length=8, max_length=255)],
) -> AiJob:
    workspace_id, user_id, api_key_id, request_id = _identity(request)
    return await request.app.state.container.ai.create(
        workspace_id=workspace_id,
        user_id=user_id,
        api_key_id=api_key_id,
        request_id=request_id,
        payload=payload,
        idempotency_key=idempotency_key,
    )


@router.get("/jobs/{generation_id}", response_model=AiJob)
async def get_job(
    request: Request,
    generation_id: Annotated[str, Path(min_length=36, max_length=36)],
) -> AiJob:
    workspace_id, user_id, _, _ = _identity(request)
    return await request.app.state.container.ai.get(
        workspace_id=workspace_id,
        user_id=user_id,
        generation_id=generation_id,
    )


@router.post("/jobs/{generation_id}/cancel", response_model=AiJob)
async def cancel_job(
    request: Request,
    generation_id: Annotated[str, Path(min_length=36, max_length=36)],
) -> AiJob:
    workspace_id, user_id, api_key_id, request_id = _identity(request)
    return await request.app.state.container.ai.cancel(
        workspace_id=workspace_id,
        user_id=user_id,
        api_key_id=api_key_id,
        request_id=request_id,
        generation_id=generation_id,
    )