Tsah00/sql-env

1"""2server/sql_environment.py - Core SQL Query Learning Environment.3 4Implements the OpenEnv Environment interface with:5  - reset(): Initialize a fresh episode with a SQLite database6  - step(action): Execute an agent SQL query, grade it, return observation7  - state: Current episode metadata8 9The environment hosts an in-memory SQLite database seeded with e-commerce10data across 4 tables (customers, products, orders, order_items).11"""12 13from __future__ import annotations14 15import sqlite316import uuid17from typing import Optional18 19from models import SQLAction, SQLObservation, SQLState20from server.tasks import TASKS, SCHEMA_INFO, grade, seed_database21 22 23# Max steps per episode before forced termination24MAX_STEPS = 2025 26 27class SQLEnvironment:28    """29    SQL Query Learning Environment.30 31    At each step the agent submits a SQL query string. The environment:32      1. Executes the query against the local SQLite database.33      2. Compares the result against the reference answer for the active task.34      3. Returns an observation with the result, reward, and feedback.35 36    Episode flow:37      reset(difficulty='easy'|'medium'|'hard') -> SQLObservation38      step(SQLAction)                          -> SQLObservation39      state                                    -> SQLState40    """41 42    def __init__(self) -> None:43        self._conn: Optional[sqlite3.Connection] = None44        self._state = SQLState()45        self._current_task: Optional[dict] = None46        self._task_ids_by_difficulty = {47            "easy":   ["easy_1", "easy_2", "easy_3"],48            "medium": ["medium_1", "medium_2", "medium_3"],49            "hard":   ["hard_1", "hard_2", "hard_3"],50        }51        self._task_index: int = 052 53    # ------------------------------------------------------------------54    # Public API55    # ------------------------------------------------------------------56 57    def reset(self, difficulty: str = "easy", task_id: Optional[str] = None) -> SQLObservation:58        """59        Start a new episode.60 61        Args:62            difficulty: Task difficulty tier ('easy', 'medium', 'hard').63            task_id: Optional specific task ID to start with.64 65        Returns:66            Initial SQLObservation with schema info and task description.67        """68        # Fresh in-memory SQLite database for each episode69        if self._conn:70            self._conn.close()71        self._conn = sqlite3.connect(":memory:")72        self._conn.row_factory = sqlite3.Row73        seed_database(self._conn)74 75        # Select the first task for the given difficulty76        if task_id and task_id in TASKS:77            self._current_task = TASKS[task_id]78        else:79            ids = self._task_ids_by_difficulty.get(difficulty, ["easy_1"])80            self._task_index = 081            self._current_task = TASKS[ids[0]]82 83        # Reset state84        self._state = SQLState(85            episode_id=str(uuid.uuid4()),86            step_count=0,87            current_task_id=self._current_task["id"],88            current_difficulty=self._current_task["difficulty"],89            total_reward=0.0,90            tasks_completed=0,91            best_reward=0.0,92            last_query="",93            last_error="",94        )95 96        return SQLObservation(97            result=[],98            error="",99            reward=0.0,100            done=False,101            message=(102                "Episode started. Write a SQL query to complete the task.\n"103                "You can explore the schema using: "104                "SELECT name FROM sqlite_master WHERE type='table';"105            ),106            schema_info=SCHEMA_INFO,107            task_description=self._current_task["description"],108            expected_columns=self._current_task["expected_columns"],109            step_count=0,110            score_breakdown={},111        )112 113    def step(self, action: SQLAction) -> SQLObservation:114        """115        Execute a SQL query and return graded observation.116 117        Args:118            action: SQLAction with the query string to execute.119 120        Returns:121            SQLObservation with result rows, reward, and feedback.122        """123        if self._conn is None or self._current_task is None:124            return SQLObservation(125                error="Environment not initialised. Call reset() first.",126                reward=0.0,127                done=True,128                message="Call reset() to start a new episode.",129            )130 131        self._state.step_count += 1132        self._state.last_query = action.query133 134        # Determine which task to grade135        task_id = action.task_id or self._state.current_task_id136 137        # Override difficulty if requested138        if action.difficulty and action.difficulty != self._state.current_difficulty:139            ids = self._task_ids_by_difficulty.get(action.difficulty, ["easy_1"])140            task_id = ids[0]141            self._current_task = TASKS.get(task_id, self._current_task)142            self._state.current_task_id = task_id143            self._state.current_difficulty = action.difficulty144 145        # Grade the query146        reward, message, agent_rows, expected_rows = grade(147            task_id, action.query, self._conn148        )149 150        # Update state151        self._state.total_reward += reward152        if reward > self._state.best_reward:153            self._state.best_reward = reward154 155        # Perfect score: advance to next task (if any)156        done = False157        if reward >= 1.0:158            self._state.tasks_completed += 1159            message += " Moving to next task..."160            advanced = self._advance_task()161            if not advanced:162                done = True163                message = "All tasks completed! Episode done."164 165        # Force done after max steps166        if self._state.step_count >= MAX_STEPS:167            done = True168            message += f" Episode ended (max {MAX_STEPS} steps reached)."169 170        self._state.last_error = ""171        score_breakdown = {172            "correctness": round(reward * 0.7 / max(reward, 0.001), 4) if reward > 0 else 0.0,173            "keyword_bonus": 0.1,174            "efficiency_bonus": round(reward - min(reward * 0.7, 0.7), 4),175        }176 177        return SQLObservation(178            result=agent_rows,179            error="",180            reward=reward,181            done=done,182            message=message,183            schema_info=SCHEMA_INFO,184            task_description=self._current_task["description"] if not done else "Episode complete.",185            expected_columns=self._current_task["expected_columns"] if not done else [],186            step_count=self._state.step_count,187            score_breakdown=score_breakdown,188        )189 190    @property191    def state(self) -> SQLState:192        """Return current episode state."""193        return self._state194 195    # ------------------------------------------------------------------196    # Private helpers197    # ------------------------------------------------------------------198 199    def _advance_task(self) -> bool:200        """201        Move to the next task in the current difficulty tier.202 203        Returns True if advanced, False if no more tasks remain.204        """205        difficulty = self._state.current_difficulty206        ids = self._task_ids_by_difficulty.get(difficulty, [])207        self._task_index += 1208        if self._task_index < len(ids):209            next_id = ids[self._task_index]210            self._current_task = TASKS[next_id]211            self._state.current_task_id = next_id212            return True213        return False214 215    def close(self) -> None:216        """Close the SQLite connection."""217        if self._conn:218            self._conn.close()219            self._conn = None220