"""Image data models."""
from datetime import datetime
from pydantic import BaseModel, Field
class GeneratedImage(BaseModel):
"""Represents a generated image from the API."""
# Image data
b64_json: str | None = Field(
default=None, description="Base64-encoded image data (PNG format)"
)
url: str | None = Field(default=None, description="URL to the generated image")
# Metadata
revised_prompt: str | None = Field(
default=None, description="The prompt as revised by the model"
)
created_at: datetime = Field(
default_factory=datetime.now, description="When the image was generated"
)
# Generation parameters (for reference)
model: str | None = Field(default=None, description="Model used for generation")
size: str | None = Field(default=None, description="Image size")
quality: str | None = Field(default=None, description="Image quality setting")
style: str | None = Field(default=None, description="Image style setting")
@property
def has_data(self) -> bool:
"""Check if the image has data (either base64 or URL)."""
return self.b64_json is not None or self.url is not None
class ImageGenerationError(BaseModel):
"""Represents an error during image generation."""
index: int = Field(description="Index of the failed generation attempt")
error: str = Field(description="Error message")
error_type: str | None = Field(default=None, description="Type of error")
timestamp: datetime = Field(
default_factory=datetime.now, description="When the error occurred"
)
