Spaces:
Sleeping
Sleeping
File size: 5,637 Bytes
be62603 f38af3b be62603 f38af3b be62603 f38af3b be62603 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """
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,
)
|