src/kittykat_agent/categorization/models.py

from typing import List
from pydantic import BaseModel, Field, field_validator


class ImageItem(BaseModel):
    id: str
    asset_url: str


class CategoryGroup(BaseModel):
    category_name: str
    images: List[ImageItem]


class CategorizationRequest(BaseModel):
    brand_id: str = Field(..., description="Brand ID to filter gallery items")
    categories: List[str] = Field(..., min_length=3,
                                  description="List of category names (minimum 3)")

    @field_validator('categories')
    def validate_categories(cls, v):
        if len(v) < 3:
            raise ValueError('At least 3 categories are required')
        # Remove duplicates while preserving order
        seen = set()
        unique_categories = []
        for cat in v:
            if cat not in seen:
                seen.add(cat)
                unique_categories.append(cat)
        return unique_categories


class CategorizationResponse(BaseModel):
    total_images: int
    total_categories: int
    categories: List[CategoryGroup]