YG-Datasets/backend/app/schemas/base.py

"""
Base Pydantic schemas
"""
from datetime import datetime
from typing import Optional, Any
from uuid import UUID
from pydantic import BaseModel, ConfigDict


class TimestampMixin(BaseModel):
    """Mixin for timestamps"""
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None


class UUIDMixin(BaseModel):
    """Mixin for UUID"""
    model_config = ConfigDict(from_attributes=True)

    id: UUID


class ProjectBase(BaseModel):
    """Base project schema"""
    name: str
    description: Optional[str] = None


class ProjectCreate(ProjectBase):
    """Project create schema"""
    pass


class ProjectUpdate(ProjectBase):
    """Project update schema"""
    pass


class ProjectResponse(ProjectBase, UUIDMixin, TimestampMixin):
    """Project response schema"""
    pass


class FileBase(BaseModel):
    """Base file schema"""
    filename: str
    file_type: str
    size: Optional[int] = None


class FileResponse(FileBase, UUIDMixin, TimestampMixin):
    """File response schema"""
    status: str


class ChunkBase(BaseModel):
    """Base chunk schema"""
    name: Optional[str] = None
    content: str
    summary: Optional[str] = None
    word_count: Optional[int] = None


class ChunkCreate(ChunkBase):
    """Chunk create schema"""
    file_id: Optional[UUID] = None


class ChunkResponse(ChunkBase, UUIDMixin, TimestampMixin):
    """Chunk response schema"""
    pass


class QuestionBase(BaseModel):
    """Base question schema"""
    content: str
    answer: Optional[str] = None
    question_type: Optional[str] = None


class QuestionCreate(QuestionBase):
    """Question create schema"""
    chunk_id: Optional[UUID] = None


class QuestionResponse(QuestionBase, UUIDMixin, TimestampMixin):
    """Question response schema"""
    source: str


class DatasetBase(BaseModel):
    """Base dataset schema"""
    name: str
    description: Optional[str] = None
    dataset_type: Optional[str] = None


class DatasetCreate(DatasetBase):
    """Dataset create schema"""
    pass


class DatasetResponse(DatasetBase, UUIDMixin, TimestampMixin):
    """Dataset response schema"""
    question_count: Optional[int] = None


class EvalDatasetBase(BaseModel):
    """Base eval dataset schema"""
    name: str
    question_type: Optional[str] = None


class EvalDatasetCreate(EvalDatasetBase):
    """Eval dataset create schema"""
    pass


class EvalDatasetResponse(EvalDatasetBase, UUIDMixin, TimestampMixin):
    """Eval dataset response schema"""
    pass


class TagBase(BaseModel):
    """Base tag schema"""
    label: str
    parent_id: Optional[UUID] = None
    color: Optional[str] = None


class TagCreate(TagBase):
    """Tag create schema"""
    pass


class TagResponse(TagBase, UUIDMixin, TimestampMixin):
    """Tag response schema"""
    pass


class ModelConfigBase(BaseModel):
    """Base model config schema"""
    provider: str
    model_name: Optional[str] = None
    api_key: Optional[str] = None
    api_base: Optional[str] = None
    is_default: Optional[str] = "false"


class ModelConfigCreate(ModelConfigBase):
    """Model config create schema"""
    pass


class ModelConfigResponse(ModelConfigBase, UUIDMixin, TimestampMixin):
    """Model config response schema"""
    pass


class TaskBase(BaseModel):
    """Base task schema"""
    task_type: str
    status: Optional[str] = "pending"
    progress: Optional[int] = 0


class TaskResponse(TaskBase, UUIDMixin, TimestampMixin):
    """Task response schema"""
    result: Optional[Any] = None
    error: Optional[str] = None