app/schemas/requests.py

from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel

from app.config import get_settings

settings = get_settings()


class Attribute(BaseModel):
    description: str
    data_type: str
    allowed_values: Optional[List[str]] = []


class ExtractionRequest(BaseModel):
    attributes: Dict[str, Attribute]
    img_urls: Optional[List[str]] = None
    product_taxonomy: str
    request_meta: Optional[Dict[str, str]] = None
    product_data: Optional[Dict[str, Union[str, List[str]]]] = None
    ai_model: str = settings.DEFAULT_MODEL  # type: ignore
    max_attempts: int = settings.DEFAULT_MAX_ATTEMPTS  # type: ignore


class FollowSchemaRequest(BaseModel):
    data_schema: Dict[str, Any]
    data: Dict[str, Any]
    request_meta: Optional[Dict[str, str]] = None
    ai_model: str = settings.DEFAULT_MODEL
    max_attempts: int = settings.DEFAULT_MAX_ATTEMPTS  # type: ignore


class ResultRequest(BaseModel):
    task_id: str