SQL-Detective-v1

from openenv.core import Environment, StepResult

from .models import SQLAction, SQLObservation, SQLState
from .sandbox import ReadOnlySQLite
from .grader import grade


class SQLDetectiveEnv(Environment[SQLAction, SQLObservation, SQLState]):
    def __init__(self, split="train", max_turns=12, query_timeout_s=5, row_limit=200):
        self.questions = load_questions(split)
        self.max_turns = max_turns
        self.timeout = query_timeout_s
        self.row_limit = row_limit

    def reset(self, seed=None) -> SQLObservation:
        self.task = self.questions.sample(seed)
        self.db = ReadOnlySQLite(self.task.database, timeout_s=self.timeout)
        self.state = SQLState(turn=0, invalid_queries=0)
        return SQLObservation.initial(self.task.question)

    def step(self, action: SQLAction) -> StepResult[SQLObservation]:
        self.state.turn += 1
        if action.tool == "submit_answer":
            reward = grade(action.args["rows"], self.task.expected, self.state)
            return StepResult(SQLObservation.done(), reward=reward, terminated=True)

        result = self.db.call(action.tool, action.args, row_limit=self.row_limit)
        if result.error:
            self.state.invalid_queries += 1
        truncated = self.state.turn >= self.max_turns
        return StepResult(SQLObservation.tool(result), reward=0.0, truncated=truncated)