models.py

from typing import Optional, Dict, Any
from pydantic import Field
from openenv.core.env_server.types import Action as BaseAction, Observation as BaseObservation


class Observation(BaseObservation):
    task_id: int = Field(default=0, description="The ID of the task to perform.")
    query: str = Field(default="", description="The SQL query to review and optimize.")
    schema_context: str = Field(default="", description="The database schema context.")
    hint: Optional[str] = Field(default=None, description="An optional natural-language hint.")
    step_number: int = Field(default=0, description="The current step number in the episode.")
    max_steps: int = Field(default=0, description="The maximum allowed steps for this task.")


class Action(BaseAction):
    rewritten_query: str = Field(default="", description="The rewritten, optimized SQL query.")
    explanation: str = Field(default="", description="A brief explanation of the changes.")
    is_done: bool = Field(default=False, description="Set to true to submit for final scoring.")


class Reward:
    def __init__(self, score: float = 0.0, breakdown: Dict[str, float] = None, feedback: str = ""):
        self.score = score
        self.breakdown = breakdown or {}
        self.feedback = feedback

    def model_dump(self):
        return {"score": self.score, "breakdown": self.breakdown, "feedback": self.feedback}