Spaces:
Sleeping
Sleeping
File size: 4,091 Bytes
e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 2052f0b e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 2052f0b e386bcf 4616824 e386bcf 4616824 e386bcf 4616824 e386bcf 2052f0b e386bcf | 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 | 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)
|