models.py

from pydantic import BaseModel, Field
from typing import Optional
from enum import Enum


class OpenRouterModel(str, Enum):
    GLM_5 = "z-ai/glm-5"
    GPT_4O_MINI = "openai/gpt-4o-mini"
    CLAUDE_SONNET = "anthropic/claude-sonnet-4.6"


class ColumnConfig(BaseModel):
    name: str = Field(..., description="Column name")
    type: str = Field(..., description="Column type: sampler, llm_text")
    params: dict = Field(default_factory=dict)


class GenerateRequest(BaseModel):
    num_records: int = Field(default=10, ge=1, le=100)
    model: OpenRouterModel = Field(default=OpenRouterModel.GLM_5)
    columns: list[ColumnConfig] = Field(...)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=256, ge=64, le=4096)


class PreviewRequest(BaseModel):
    model: OpenRouterModel = Field(default=OpenRouterModel.GLM_5)
    columns: list[ColumnConfig] = Field(...)
    temperature: float = Field(default=0.7)
    max_tokens: int = Field(default=256)


class GenerateResponse(BaseModel):
    success: bool
    data: Optional[list[dict]] = None
    record_count: int = 0
    error: Optional[str] = None


class PreviewResponse(BaseModel):
    success: bool
    sample: Optional[dict] = None
    error: Optional[str] = None


class HealthResponse(BaseModel):
    status: str
    model: str
    api_configured: bool