Spaces:
Sleeping
Sleeping
File size: 4,862 Bytes
5716a3a | 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 | import random
from typing import Optional
from sql_env.models import (
SQLObservation, SQLAction, SQLTask, StepResult
)
from sql_env.grader import grade, generate_feedback
from sql_env.tasks import ALL_TASKS
class SQLCorrectionEnv:
"""
OpenEnv-compliant SQL Query Correction Environment.
The agent receives a broken SQL query and must return the corrected version.
Reward is shaped across the full trajectory — partial credit is given for
incremental improvements, penalizing stagnation and infinite loops.
"""
def __init__(self, difficulty: str = "easy", task_index: Optional[int] = None):
if difficulty not in ALL_TASKS:
raise ValueError(f"difficulty must be one of {list(ALL_TASKS.keys())}")
self.difficulty = difficulty
self.task_index = task_index
self._task: Optional[SQLTask] = None
self._step_count: int = 0
self._done: bool = False
self._previous_attempt: Optional[str] = None
self._last_feedback: Optional[str] = None
self._last_reward: float = 0.0
self._stagnation_count: int = 0
# ── OpenEnv Interface ─────────────────────────────────────
async def reset(self) -> SQLObservation:
"""Reset the environment and return the initial observation."""
tasks = ALL_TASKS[self.difficulty]
if self.task_index is not None:
self._task = tasks[self.task_index % len(tasks)]
else:
self._task = random.choice(tasks)
self._step_count = 0
self._done = False
self._previous_attempt = None
self._last_feedback = None
self._last_reward = 0.0
self._stagnation_count = 0
return self._make_observation()
async def step(self, action: SQLAction) -> StepResult:
"""
Take one step: grade the agent's corrected query and return
(observation, reward, done, info).
"""
if self._done:
raise RuntimeError("Episode is done. Call reset() to start a new episode.")
if self._task is None:
raise RuntimeError("Environment not initialized. Call reset() first.")
self._step_count += 1
# grade the action
reward_model = grade(action, self._task)
reward = reward_model.value
# detect stagnation (same reward twice in a row) — penalize
if abs(reward - self._last_reward) < 0.01 and self._step_count > 1:
self._stagnation_count += 1
if self._stagnation_count >= 2:
reward = max(0.0, reward - 0.1) # stagnation penalty
else:
self._stagnation_count = 0
self._last_reward = reward
# generate feedback for the next observation
feedback = generate_feedback(action, self._task, reward_model)
self._last_feedback = feedback
self._previous_attempt = action.corrected_query
# episode ends on perfect score or max steps reached
done = reward_model.value == 1.0 or self._step_count >= self._task.max_steps
self._done = done
obs = self._make_observation()
return StepResult(
observation=obs,
reward=round(reward, 4),
done=done,
info={
"grader_reason": reward_model.reason,
"step": self._step_count,
"max_steps": self._task.max_steps,
"task_id": self._task.task_id,
}
)
async def state(self) -> dict:
"""Return the current internal state of the environment."""
if self._task is None:
return {"status": "not_initialized"}
return {
"task_id": self._task.task_id,
"difficulty": self.difficulty,
"step_count": self._step_count,
"done": self._done,
"last_reward": self._last_reward,
"max_steps": self._task.max_steps,
"previous_attempt": self._previous_attempt,
}
async def close(self):
"""Clean up resources."""
self._task = None
self._done = True
# ── Internal ──────────────────────────────────────────────
def _make_observation(self) -> SQLObservation:
assert self._task is not None
return SQLObservation(
task_id=self._task.task_id,
broken_query=self._task.broken_query,
schema_context=self._task.schema_context,
error_hint=self._task.error_hint if self.difficulty == "easy" else None,
step_number=self._step_count,
previous_attempt=self._previous_attempt,
feedback=self._last_feedback,
)
|