File size: 2,720 Bytes
3493993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
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,
    )