src/api/models/embedding_models.py

from pydantic import BaseModel
from typing import List, Dict


# Pydantic models for request validation
class CreateEmbeddingRequest(BaseModel):
    query: str
    target_column: str = "product_type"
    output_column: str = "embedding"
    model: str = "text-embedding-3-small"
    batch_size: int = 10
    max_concurrent_requests: int = 10
    dataset_name: str = "re-mind/product_type_embedding"


class UpdateEmbeddingRequest(BaseModel):
    dataset_name: str
    updates: Dict[str, List]  # Column name -> List of values


class DeleteEmbeddingRequest(BaseModel):
    dataset_name: str