Spaces:
Sleeping
Sleeping
| import sqlite3 | |
| import time | |
| from typing import Optional, Dict, Any, Tuple, List | |
| from .models import SQLObservation, SQLAction, SQLReward | |
| from .tasks import TASKS, load_fixtures | |
| from .graders import grade_sql | |
| class SQLReviewEnv: | |
| def __init__(self): | |
| self.conn = None | |
| self.task_id = "syntax-fix" | |
| self.step_count = 0 | |
| self.max_steps = 8 | |
| self.last_reward = 0.01 | |
| self.done = False | |
| self.last_error: Optional[str] = None | |
| self._query_start_time = 0 | |
| self.history: List[float] = [] | |
| async def reset(self, task_id: str = "syntax-fix") -> SQLObservation: | |
| self.task_id = task_id | |
| if self.task_id not in TASKS: | |
| self.task_id = "syntax-fix" | |
| self.step_count = 0 | |
| self.done = False | |
| self.last_reward = 0.01 | |
| self.last_error = None | |
| self.history = [0.01] | |
| if self.conn: | |
| self.conn.close() | |
| self.conn = sqlite3.connect(":memory:") | |
| self.conn.row_factory = sqlite3.Row | |
| # Absolute execution speed limits and RAM constraints | |
| self.conn.execute("PRAGMA max_page_count = 10000;") | |
| def timeout_handler(): | |
| if time.time() - self._query_start_time > 1.5: | |
| # Return 1 to violently abort the massive sqlite block | |
| return 1 | |
| return 0 | |
| self.conn.set_progress_handler(timeout_handler, 1000) | |
| # Anchor start time before fixtures load so the timeout is valid immediately | |
| self._query_start_time = time.time() | |
| load_fixtures(self.conn, self.task_id) | |
| task_data = TASKS[self.task_id] | |
| return SQLObservation( | |
| task_id=self.task_id, | |
| db_schema=task_data["db_schema"], | |
| query=task_data["query"], | |
| error_message=None, | |
| expected_hint=task_data["expected_hint"], | |
| step=self.step_count | |
| ) | |
| async def step(self, action: SQLAction) -> SQLReward: | |
| if self.done: | |
| return SQLReward(value=self.last_reward, breakdown={}, done=True, info={"error": self.last_error or ""}) | |
| self.step_count += 1 | |
| # Anchor the time bounds logic prior to grader evaluation | |
| self._query_start_time = time.time() | |
| task_data = TASKS[self.task_id] | |
| expected_sql = task_data["validation_query"] | |
| reward_val, breakdown, done, info = grade_sql( | |
| self.task_id, self.conn, action.sql, expected_sql, self.step_count, self.max_steps | |
| ) | |
| self.last_reward = reward_val | |
| self.done = done | |
| self.history.append(reward_val) | |
| # Persist the latest error message for next observation if needed | |
| self.last_error = info.get("error") or info.get("validation_error") or info.get("plan_error") or None | |
| return SQLReward( | |
| value=reward_val, | |
| breakdown=breakdown, | |
| done=done, | |
| info=info | |
| ) | |
| def state(self) -> dict: | |
| return { | |
| "task_id": self.task_id, | |
| "current_step": self.step_count, | |
| "max_steps": self.max_steps, | |
| "last_reward": self.last_reward, | |
| "done": self.done, | |
| "history": self.history | |
| } | |
| def close(self) -> None: | |
| """Cleanly close the SQLite connection. Called on server shutdown.""" | |
| if self.conn: | |
| self.conn.close() | |
| self.conn = None | |