66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
"""
|
|
Pydantic models for request/response schemas
|
|
"""
|
|
from pydantic import BaseModel, Field
|
|
from typing import Optional
|
|
from enum import Enum
|
|
|
|
|
|
class TaskStatus(str, Enum):
|
|
PENDING = "pending"
|
|
PROCESSING = "processing"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
|
|
class GenerateRequest(BaseModel):
|
|
"""Request model for image generation"""
|
|
prompt: str = Field(..., description="Image generation prompt", min_length=1)
|
|
cookies: str = Field(..., description="Meta AI session cookies")
|
|
num_images: int = Field(default=4, ge=1, le=8, description="Number of images to generate")
|
|
|
|
|
|
class GrokChatRequest(BaseModel):
|
|
"""Request model for Grok chat"""
|
|
message: str = Field(..., description="Message content")
|
|
history: Optional[list] = Field(default=None, description="Chat history")
|
|
cookies: Optional[dict] = Field(default=None, description="Grok session cookies")
|
|
user_agent: Optional[str] = Field(default=None, description="Browser User-Agent")
|
|
|
|
|
|
class ImageResult(BaseModel):
|
|
"""Single generated image result"""
|
|
url: str
|
|
data: Optional[str] = None # base64 encoded image data
|
|
prompt: str
|
|
model: str = "imagine"
|
|
|
|
|
|
class GenerateResponse(BaseModel):
|
|
"""Response model for image generation"""
|
|
success: bool
|
|
images: list[ImageResult] = []
|
|
error: Optional[str] = None
|
|
task_id: Optional[str] = None
|
|
|
|
|
|
class GrokChatResponse(BaseModel):
|
|
"""Response model for Grok chat"""
|
|
response: str
|
|
error: Optional[str] = None
|
|
|
|
|
|
class TaskStatusResponse(BaseModel):
|
|
"""Response model for async task status"""
|
|
task_id: str
|
|
status: TaskStatus
|
|
images: list[ImageResult] = []
|
|
error: Optional[str] = None
|
|
progress: Optional[int] = None # 0-100
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
"""Health check response"""
|
|
status: str = "healthy"
|
|
version: str = "1.0.0"
|
|
browser_ready: bool = True
|