ShrijaSambhavi/AI-Evaluation-Environment-for-SQL-Query-Optimization

1from typing import Dict2from openenv.core import EnvClient3from openenv.core.client_types import StepResult4from openenv.core.env_server.types import State5from .models import SqlEnvAction, SqlEnvObservation6 7class SqlEnv(EnvClient[SqlEnvAction, SqlEnvObservation, State]):8 9    def _step_payload(self, action: SqlEnvAction) -> Dict:10        return {11            "action_type": action.action_type,12            "query": action.query,13        }14 15    def _parse_result(self, payload: Dict) -> StepResult[SqlEnvObservation]:16        obs_data = payload.get("observation", payload)  # some servers inline obs17        observation = SqlEnvObservation(18            task_description=obs_data.get("task_description", ""),19            schema_info=obs_data.get("schema_info", ""),20            initial_query=obs_data.get("initial_query"),21            feedback=obs_data.get("feedback", ""),22            done=payload.get("done", False),23            reward=payload.get("reward", 0.0),  # ✅ explicit default24        )25        return StepResult(26            observation=observation,27            reward=payload.get("reward", 0.0),  # ✅ never None28            done=payload.get("done", False),29        )30 31    def _parse_state(self, payload: Dict) -> State:32        return State(33            episode_id=payload.get("episode_id"),34            step_count=payload.get("step_count", 0),35        )