app/api/generation.py

from __future__ import annotations

from typing import Annotated
from uuid import UUID

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

from app.generation.schemas.requests import (
    GenerationJobAttemptView,
    GenerationJobView,
    GenerationProviderView,
    GenerationRequestCreate,
    GenerationRequestView,
)
from app.generation.model_registry import GenerationModelView

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


def _generation(request: Request):
    service = request.app.state.container.generation
    service.ensure_ready()
    return service


def _identity(request: Request) -> tuple[str, str]:
    context = request.state.auth
    if not context.workspace_id or not context.user_id:
        # The credential itself is never a workspace.  API-key middleware
        # resolves the authoritative membership before this route runs.
        raise HTTPException(status_code=403, detail="No active workspace membership.")
    return context.workspace_id, context.user_id


@router.get("/providers", response_model=list[GenerationProviderView])
async def list_providers(request: Request) -> list[GenerationProviderView]:
    return _generation(request).list_providers()


@router.get(
    "/providers/{provider}", response_model=GenerationProviderView
)
async def get_provider(request: Request, provider: str) -> GenerationProviderView:
    return _generation(request).get_provider(provider)


@router.get("/models", response_model=list[GenerationModelView])
async def list_models(
    request: Request,
    provider: str | None = Query(default=None, min_length=1, max_length=64),
) -> list[GenerationModelView]:
    return _generation(request).list_models(provider=provider)


@router.get("/providers/{provider}/models/{model_id}", response_model=GenerationModelView)
async def get_model(
    request: Request, provider: str, model_id: str
) -> GenerationModelView:
    return _generation(request).get_model(provider, model_id)


@router.post(
    "/requests",
    response_model=GenerationRequestView,
    status_code=status.HTTP_202_ACCEPTED,
)
async def create_request(
    request: Request,
    payload: GenerationRequestCreate,
    idempotency_key: Annotated[
        str, Header(alias="Idempotency-Key", min_length=8, max_length=255)
    ],
) -> GenerationRequestView:
    workspace_id, user_id = _identity(request)
    return await _generation(request).create(
        workspace_id=workspace_id,
        user_id=user_id,
        payload=payload,
        idempotency_key=idempotency_key,
        project_id=payload.project_id,
    )


@router.get("/requests", response_model=list[GenerationRequestView])
async def list_requests(
    request: Request,
    offset: int = Query(default=0, ge=0),
    limit: int = Query(default=100, ge=1, le=500),
    project_id: UUID | None = Query(default=None),
) -> list[GenerationRequestView]:
    workspace_id, user_id = _identity(request)
    return await _generation(request).list_requests(
        workspace_id,
        user_id,
        offset=offset,
        limit=limit,
        project_id=str(project_id) if project_id else None,
    )


@router.get("/requests/{generation_request_id}", response_model=GenerationRequestView)
async def get_request(
    request: Request, generation_request_id: str
) -> GenerationRequestView:
    workspace_id, user_id = _identity(request)
    return await _generation(request).get_request(
        workspace_id, user_id, generation_request_id
    )


@router.get("/jobs/{generation_job_id}", response_model=GenerationJobView)
async def get_job(request: Request, generation_job_id: str) -> GenerationJobView:
    workspace_id, user_id = _identity(request)
    return await _generation(request).get_job(
        workspace_id, user_id, generation_job_id
    )


@router.post(
    "/jobs/{generation_job_id}/retry",
    response_model=GenerationRequestView,
    status_code=status.HTTP_202_ACCEPTED,
)
async def retry_failed_job(
    request: Request,
    generation_job_id: str,
    idempotency_key: Annotated[
        str, Header(alias="Idempotency-Key", min_length=8, max_length=255)
    ],
) -> GenerationRequestView:
    workspace_id, user_id = _identity(request)
    return await _generation(request).retry_failed_job(
        workspace_id,
        user_id,
        generation_job_id,
        idempotency_key=idempotency_key,
    )


@router.get(
    "/jobs/{generation_job_id}/attempts",
    response_model=list[GenerationJobAttemptView],
)
async def list_job_attempts(
    request: Request, generation_job_id: str
) -> list[GenerationJobAttemptView]:
    workspace_id, user_id = _identity(request)
    return await _generation(request).list_job_attempts(
        workspace_id, user_id, generation_job_id
    )


@router.post(
    "/jobs/{generation_job_id}/cancel", response_model=GenerationJobView
)
async def cancel_job(request: Request, generation_job_id: str) -> GenerationJobView:
    workspace_id, user_id = _identity(request)
    return await _generation(request).cancel(
        workspace_id, user_id, generation_job_id
    )