server/environment.py

import sqlite3
import uuid

from openenv.core.env_server import Environment
from models import SQLAction, SQLObservation, SQLState
from server.tasks import TASKS


class SQLDebuggerEnv(Environment):
    # Class-level state — shared across instances because the framework
    # creates a new instance for every HTTP request (reset and step).
    _db = None
    _current_task = None
    _task_index = 0
    _state = SQLState()

    def reset(self) -> SQLObservation:
        # 1. Pick the next task (cycles through 0..8..0..8..)
        cls = type(self)
        cls._current_task = TASKS[cls._task_index]
        cls._task_index = (cls._task_index + 1) % len(TASKS)

        # 2. Close old DB if any, create a fresh one
        if cls._db:
            cls._db.close()
        cls._db = sqlite3.connect(":memory:")

        # 3. Set up tables and insert seed data
        for stmt in cls._current_task["schema"]:
            cls._db.execute(stmt)
        for stmt in cls._current_task["seed_data"]:
            cls._db.execute(stmt)
        cls._db.commit()

        # 4. Update episode metadata
        cls._state = SQLState(
            task_id=cls._current_task["id"],
            difficulty=cls._current_task["difficulty"],
            step_count=0,
            episode_id=str(uuid.uuid4()),
        )

        # 5. Return the puzzle
        return SQLObservation(
            broken_query=cls._current_task["broken_query"],
            table_schema=cls._current_task["schema"],
            description=cls._current_task["description"],
            expected_output=cls._current_task["expected_output"],
        )

    def step(self, action: SQLAction) -> SQLObservation:
        cls = type(self)

        # 1. Track step count
        cls._state.step_count += 1

        # 2. Score the agent's attempted fix
        score, actual_output, error_message = self.compute_reward(action.corrected_query)

        # 3. One fewer attempt remaining
        attempts_left = 3 - cls._state.step_count

        # 4. Done if perfect score or no attempts left
        is_done = score >= 0.99 or attempts_left <= 0

        # 5. Return updated observation
        return SQLObservation(
            broken_query=cls._current_task["broken_query"],
            table_schema=cls._current_task["schema"],
            description=cls._current_task["description"],
            expected_output=cls._current_task["expected_output"],
            actual_output=actual_output,
            score=score,
            error_message=error_message,
            done=is_done,
            reward=score,
            attempts_left=attempts_left,
        )

    @property
    def state(self) -> SQLState:
        return type(self)._state

    def compute_reward(self, submitted_sql: str) -> tuple:
        """
        Compute partial credit score for the submitted SQL.
        Returns (score, actual_output, error_message).
        """
        cls = type(self)
        score = 0.05
        actual_output = None
        error_message = None
        expected = cls._current_task["expected_output"]

        # Level 1: Does it parse?
        try:
            cls._db.execute("EXPLAIN " + submitted_sql)
            score = 0.2
        except Exception as e:
            error_message = str(e)
            return (score, actual_output, error_message)

        # Level 2: Does it run without crashing?
        try:
            actual_output = cls._db.execute(submitted_sql).fetchall()
            score = 0.4
        except Exception as e:
            error_message = str(e)
            return (score, actual_output, error_message)

        # Level 3: Right number of columns?
        if actual_output and expected and len(actual_output[0]) == len(expected[0]):
            score = 0.6

            # Level 4: Right number of rows?
            if len(actual_output) == len(expected):
                score = 0.8

                # Level 5: Exact match? (sort to ignore row order)
                if sorted(actual_output) == sorted(expected):
                    score = 0.99

        return (score, actual_output, error_message)