Spaces:
Sleeping
Sleeping
| """ | |
| SQL Debug Environment — server-side environment logic. | |
| Implements the OpenEnv Environment interface: reset(), step(), state property. | |
| """ | |
| import sys | |
| import os | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| import random | |
| from uuid import uuid4 | |
| from typing import Optional | |
| from models import Action, Observation, State | |
| from tasks import TASKS, TASK_GROUPS, get_task | |
| from graders import grade_submission | |
| MAX_STEPS = 5 | |
| class SQLDebugEnvironment: | |
| """ | |
| Real-world OpenEnv environment for SQL query debugging. | |
| An agent is given a broken SQL query and must submit corrected versions | |
| to earn rewards. The environment tracks attempts and computes shaped rewards. | |
| """ | |
| def __init__(self): | |
| self._state: Optional[State] = None | |
| self._current_task_id: Optional[str] = None | |
| self._current_task: Optional[dict] = None | |
| self._submissions: list = [] | |
| self._done: bool = False | |
| self._last_reward: float = 0.0 | |
| self._last_result: str = "" | |
| self._task_group: Optional[str] = None | |
| def reset( | |
| self, | |
| task_group: Optional[str] = None, | |
| task_id: Optional[str] = None, | |
| seed: Optional[int] = None, | |
| ) -> Observation: | |
| """ | |
| Initialize a new episode. | |
| Args: | |
| task_group: One of 'easy_syntax_fix', 'medium_logic_fix', 'hard_optimization_fix' | |
| task_id: Specific task id (overrides task_group) | |
| seed: Random seed for reproducibility | |
| """ | |
| if seed is not None: | |
| random.seed(seed) | |
| if task_id and task_id in TASKS: | |
| self._current_task_id = task_id | |
| self._task_group = None | |
| elif task_id and task_id in TASK_GROUPS: | |
| self._current_task_id = random.choice(TASK_GROUPS[task_id]) | |
| self._task_group = task_id | |
| elif task_group and task_group in TASK_GROUPS: | |
| self._current_task_id = random.choice(TASK_GROUPS[task_group]) | |
| self._task_group = task_group | |
| else: | |
| self._current_task_id = random.choice(list(TASKS.keys())) | |
| self._task_group = None | |
| self._current_task = get_task(self._current_task_id) | |
| self._submissions = [] | |
| self._done = False | |
| self._last_reward = 0.0 | |
| self._last_result = "" | |
| episode_id = str(uuid4()) | |
| self._state = State( | |
| episode_id=episode_id, | |
| step_count=0, | |
| task_id=self._task_group if self._task_group else self._current_task_id, | |
| max_steps=MAX_STEPS, | |
| cumulative_reward=0.0, | |
| ) | |
| return self._build_observation() | |
| def step(self, action: Action) -> Observation: | |
| """ | |
| Execute one fix attempt. | |
| Args: | |
| action: Agent's submitted SQL query + optional explanation | |
| Returns: | |
| Observation with reward, done flag, and feedback | |
| """ | |
| if self._done: | |
| return self._build_observation( | |
| message="Episode already done. Call reset()." | |
| ) | |
| if self._state is None: | |
| raise RuntimeError("Environment not initialized. Call reset() first.") | |
| submitted_query = action.query.strip() | |
| self._submissions.append(submitted_query) | |
| self._state.step_count += 1 | |
| score, feedback, details = grade_submission( | |
| self._current_task_id, submitted_query | |
| ) | |
| reward = self._compute_reward(score, self._state.step_count) | |
| self._last_reward = reward | |
| self._state.cumulative_reward += reward | |
| self._last_result = feedback | |
| if score == 1.0: | |
| self._done = True | |
| elif self._state.step_count >= MAX_STEPS: | |
| self._done = True | |
| return self._build_observation() | |
| def state(self) -> State: | |
| """Return current environment state metadata.""" | |
| if self._state is None: | |
| return State(episode_id="uninitialized", step_count=0) | |
| return self._state | |
| def _compute_reward(self, score: float, step_num: int) -> float: | |
| """ | |
| Shaped reward function with partial progress signals. | |
| - score=1.0: full reward, bonus for solving early | |
| - score=0.7: meaningful partial credit | |
| - score=0.5: small partial credit | |
| - score=0.3: tiny credit for at least running | |
| - score=0.0: penalty for wasted attempt (after first) | |
| """ | |
| if score == 1.0: | |
| early_bonus = max(0.0, (MAX_STEPS - step_num) * 0.05) | |
| return round(1.0 + early_bonus, 4) | |
| elif score >= 0.7: | |
| return 0.5 | |
| elif score >= 0.5: | |
| return 0.2 | |
| elif score >= 0.3: | |
| return 0.1 | |
| else: | |
| return -0.05 if step_num > 1 else 0.0 | |
| def _build_observation(self, message: str = "") -> Observation: | |
| task = self._current_task | |
| attempts_remaining = MAX_STEPS - (self._state.step_count if self._state else 0) | |
| return Observation( | |
| schema_sql=task["schema_sql"].strip(), | |
| sample_data=task["sample_data"].strip(), | |
| buggy_query=task["buggy_query"].strip(), | |
| error_message=task["error_message"], | |
| task_description=task["task_description"], | |
| expected_columns=task["expected_columns"], | |
| attempts_remaining=max(0, attempts_remaining), | |
| last_attempt_result=self._last_result + (" " + message if message else ""), | |
| reward=self._last_reward, | |
| done=self._done, | |
| success=self._done and self._last_reward >= 1.0, | |
| ) | |