103 lines
2.4 KiB
Python
103 lines
2.4 KiB
Python
"""Pydantic schemas for AI generation endpoints."""
|
|
from typing import Any
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: str # "user" | "assistant" | "system"
|
|
content: str
|
|
|
|
|
|
class ChatRequest(BaseModel):
|
|
model: str
|
|
messages: list[ChatMessage]
|
|
temperature: float = 0.7
|
|
max_tokens: int = 1024
|
|
|
|
|
|
class ChatResponse(BaseModel):
|
|
id: str
|
|
model: str
|
|
content: str
|
|
usage: dict[str, Any] | None = None
|
|
|
|
|
|
class ModelInfo(BaseModel):
|
|
id: str
|
|
name: str
|
|
context_length: int | None = None
|
|
pricing: dict[str, Any] | None = None
|
|
|
|
|
|
# --- Text generation ---
|
|
|
|
class TextRequest(BaseModel):
|
|
model: str
|
|
prompt: str = ""
|
|
system_prompt: str | None = None
|
|
messages: list[ChatMessage] | None = None
|
|
temperature: float = 0.7
|
|
max_tokens: int = 1024
|
|
|
|
|
|
class TextResponse(BaseModel):
|
|
id: str
|
|
model: str
|
|
content: str
|
|
usage: dict[str, Any] | None = None
|
|
|
|
|
|
# --- Image generation ---
|
|
|
|
class ImageRequest(BaseModel):
|
|
model: str
|
|
prompt: str
|
|
n: int = 1
|
|
size: str = "1024x1024"
|
|
aspect_ratio: str | None = None # e.g. "1:1", "16:9", "9:16"
|
|
image_size: str | None = None # e.g. "0.5K", "1K", "2K", "4K"
|
|
|
|
|
|
class ImageResult(BaseModel):
|
|
url: str | None = None
|
|
b64_json: str | None = None
|
|
revised_prompt: str | None = None
|
|
image_id: str | None = None # UUID of stored row in generated_images
|
|
|
|
|
|
class ImageResponse(BaseModel):
|
|
id: str
|
|
model: str
|
|
images: list[ImageResult]
|
|
|
|
|
|
# --- Video generation ---
|
|
|
|
class VideoRequest(BaseModel):
|
|
model: str
|
|
prompt: str
|
|
duration_seconds: int | None = None
|
|
aspect_ratio: str = "16:9"
|
|
resolution: str | None = None # e.g. "480p", "720p", "1080p"
|
|
|
|
|
|
class VideoFromImageRequest(BaseModel):
|
|
model: str
|
|
image_url: str
|
|
prompt: str
|
|
duration_seconds: int | None = None
|
|
aspect_ratio: str = "16:9"
|
|
resolution: str | None = None # e.g. "480p", "720p", "1080p"
|
|
|
|
|
|
class VideoResponse(BaseModel):
|
|
id: str # This is the job_id from the provider
|
|
db_id: str | None = None # This is the UUID from our generated_videos table
|
|
model: str
|
|
status: str # "queued" | "processing" | "completed" | "failed"
|
|
polling_url: str | None = None
|
|
video_urls: list[str] | None = None
|
|
video_url: str | None = None # first entry of video_urls for convenience
|
|
error: str | None = None
|
|
metadata: dict[str, Any] | None = None
|