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: str = Field(..., description="Only Pandera schema to validate the data.")