agents/models.py

"""Pydantic models for the agent system."""

from pydantic import BaseModel


class CatalogerOutput(BaseModel):
    category: str
    sub_category: str
    materials: list[str]
    colors: list[str]
    estimated_size: str
    tags: list[str]
    complexity: str


class CopywriterOutput(BaseModel):
    title: str
    short_desc: str
    long_desc: str
    captions: list[str]


class PricerOutput(BaseModel):
    suggested_price_min: int
    suggested_price_max: int
    reasoning: str
    cost_breakdown: str | None = None


class AgentTrace(BaseModel):
    agent_name: str
    input_text: str
    output_data: dict
    duration_ms: int
    model_id: str


class PipelineResult(BaseModel):
    image_description: str
    catalog: CatalogerOutput | None = None
    copy_data: CopywriterOutput | None = None
    pricing: PricerOutput | None = None
    traces: list[AgentTrace] = []
    total_duration_ms: int = 0