from pydantic import BaseModel, Field


class AiQueryRequest(BaseModel):
    query: str = Field(..., min_length=2, max_length=500)


class AiColumn(BaseModel):
    key: str
    label: str


class AiQueryResponse(BaseModel):
    title: str
    summary: str
    dataset: str
    llm_used: bool = False
    columns: list[AiColumn]
    rows: list[dict]
    row_count: int
    filename: str
    examples: list[str] = Field(default_factory=list)