sql_env/models.py

from pydantic import BaseModel
from typing import Optional, Dict, Any

class SQLObservation(BaseModel):
    task_id: str
    db_schema: str
    query: str
    error_message: Optional[str] = None
    expected_hint: Optional[str] = None
    step: int = 0

class SQLAction(BaseModel):
    sql: str

class SQLReward(BaseModel):
    value: float
    breakdown: Dict[str, float]
    done: bool
    info: Dict[str, Any]