SyncShift/sql-correction-env

1"""2FastAPI server using openenv.core base classes — required for validator.3"""4import random5from openenv.core.env_server.http_server import create_app6from openenv.core.env_server.interfaces import Environment7 8try:9    from sql_env.models import SQLAction, SQLObservation, SQLState10    from sql_env.tasks import TASK_SETS11    from sql_env.grader import grade, generate_feedback12except ImportError:13    from models import SQLAction, SQLObservation, SQLState14    from tasks import TASK_SETS15    from grader import grade, generate_feedback16 17 18class SQLCorrectionEnvironment(Environment):19    SUPPORTS_CONCURRENT_SESSIONS = True20 21    def __init__(self):22        self._difficulty = "easy"23        self._current_task = None24        self._step_count = 025        self._done = False26        self._last_reward = 0.0127        self._rewards_history = []28        self._stagnation_count = 029 30    def reset(self, seed=None, episode_id=None, **kwargs) -> SQLObservation:31        actual_difficulty = (32            kwargs.get("task_id")33            or kwargs.get("difficulty")34            or "easy"35        )36        self._difficulty = actual_difficulty37        tasks = TASK_SETS.get(actual_difficulty, TASK_SETS["easy"])38        self._current_task = random.choice(tasks)39        self._step_count = 040        self._done = False41        self._last_reward = 0.0142        self._rewards_history = []43        self._stagnation_count = 044        return self._make_observation(previous_attempt=None, feedback=None)45 46    def step(self, action: SQLAction) -> SQLObservation:47        if self._current_task is None:48            self.reset()49 50        self._step_count += 151        reward_obj = grade(action, self._current_task)52        reward = reward_obj.value53 54        # Stagnation penalty: penalize repeating the same score55        if abs(reward - self._last_reward) < 0.01 and self._step_count > 1:56            self._stagnation_count += 157            if self._stagnation_count >= 2:58                reward = max(0.01, reward - 0.1)59        else:60            self._stagnation_count = 061        # Final Clamp  62        reward = max(0.01, min(0.98, reward))63          64        self._last_reward = reward65 66        self._rewards_history.append(reward)67 68        done = (reward >= 0.95) or (69            self._step_count >= self._current_task.max_steps70        )71        self._done = done72        feedback = generate_feedback(action, self._current_task, reward)73 74        return self._make_observation(75            previous_attempt=action.corrected_query,76            feedback=feedback,77        )78 79    @property80    def state(self) -> SQLState:81        if self._current_task is None:82            return SQLState(83                task_id="none",84                difficulty="none",85                step_count=0,86                max_steps=0,87                done=False,88                last_reward=0.01,89                rewards_history=[],90            )91        return SQLState(92            task_id=self._current_task.task_id,93            difficulty=self._difficulty,94            step_count=self._step_count,95            max_steps=self._current_task.max_steps,96            done=self._done,97            last_reward=self._last_reward,98            rewards_history=self._rewards_history,99        )100 101    # ── Internal helpers ──────────────────────────────────────────────────────102 103    def _make_observation(104        self,105        previous_attempt: str | None,106        feedback: str | None,107    ) -> SQLObservation:108        assert self._current_task is not None109        max_steps = self._current_task.max_steps110        steps_remaining = max(0, max_steps - self._step_count)111        return SQLObservation(112            task_id=self._current_task.task_id,113            broken_query=self._current_task.broken_query,114            schema_context=self._current_task.schema_context,115            # Only surface the hint on easy tasks116            error_hint=(117                self._current_task.error_hint118                if self._difficulty == "easy"119                else None120            ),121            step_number=self._step_count,122            steps_remaining=steps_remaining,123            previous_attempt=previous_attempt,124            feedback=feedback,125        )126 127 128app = create_app(129    SQLCorrectionEnvironment,130    SQLAction,131    SQLObservation,132    env_name="sql-correction-env",133)134 135 136@app.get("/tasks")137async def list_tasks():138    """List available task difficulties with metadata."""139    return {140        "tasks": [141            {142                "id": "easy",143                "difficulty": "easy",144                "description": "Fix a single SQL keyword typo. Error hint provided.",145                "max_steps": 5,146                "count": 15,147            },148            {149                "id": "medium",150                "difficulty": "medium",151                "description": (152                    "Fix multiple errors across keywords and clauses. No hint."153                ),154                "max_steps": 5,155                "count": 15,156            },157            {158                "id": "hard",159                "difficulty": "hard",160                "description": (161                    "Fix complex multi-join queries including column name errors. "162                    "Schema provided, no hint."163                ),164                "max_steps": 4,165                "count": 10,166            },167        ]168    }169 170 171def main():172    import uvicorn173    uvicorn.run(app, host="0.0.0.0", port=7860)174 175 176if __name__ == "__main__":177    main()178