src/models.py

from pydantic import BaseModel, Field


class SQLQueryModel(BaseModel):
    sql_query: str = Field(..., description="SQL query to execute.")
    explanation: str = Field(..., description="Short explanation of the SQL query.")


class PanderaSchemaModel(BaseModel):
    schema_name: str = Field(
        ..., description="Only Pandera schema to validate the data."
    )