server/sql_query_debugger_environment.py

import sqlite3
import json
import random
from uuid import uuid4
from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..models import SqlQueryDebuggerAction, SqlQueryDebuggerObservation
except ImportError:
    from models import SqlQueryDebuggerAction, SqlQueryDebuggerObservation


SCENARIOS = [
    # ── EASY: syntax fixes ──────────────────────────────────────────
    {
        "id": "easy_1",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE employees (id INTEGER, name TEXT, salary REAL, dept TEXT);",
        "setup": [
            "INSERT INTO employees VALUES (1,'Alice',75000,'Engineering');",
            "INSERT INTO employees VALUES (2,'Bob',50000,'Marketing');",
            "INSERT INTO employees VALUES (3,'Carol',90000,'Engineering');",
        ],
        "broken_query": "SELCT name, salary FROM employees WHERE salary > 60000;",
        "fixed_query":  "SELECT name, salary FROM employees WHERE salary > 60000;",
        "error_message": "ParseError: near 'SELCT': syntax error",
        "expected_output_hint": "Should return Alice and Carol who earn more than 60000",
        "expected_rows": [("Alice", 75000.0), ("Carol", 90000.0)],
    },
    {
        "id": "easy_2",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE products (id INTEGER, name TEXT, price REAL, stock INTEGER);",
        "setup": [
            "INSERT INTO products VALUES (1,'Laptop',999.99,10);",
            "INSERT INTO products VALUES (2,'Mouse',29.99,50);",
            "INSERT INTO products VALUES (3,'Keyboard',79.99,30);",
        ],
        "broken_query": "SELECT name price FROM products WHERE stock > 20;",
        "fixed_query":  "SELECT name, price FROM products WHERE stock > 20;",
        "error_message": "",
        "expected_output_hint": "Should return Mouse and Keyboard with their prices",
        "expected_rows": [("Mouse", 29.99), ("Keyboard", 79.99)],
    },
    {
        "id": "easy_3",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE students (id INTEGER, name TEXT, grade INTEGER, subject TEXT);",
        "setup": [
            "INSERT INTO students VALUES (1,'Dan',85,'Math');",
            "INSERT INTO students VALUES (2,'Eve',92,'Science');",
            "INSERT INTO students VALUES (3,'Frank',78,'Math');",
        ],
        "broken_query": "SELECT name, grade FROM students WERE subject = 'Math';",
        "fixed_query":  "SELECT name, grade FROM students WHERE subject = 'Math';",
        "error_message": "ParseError: near 'WERE': syntax error",
        "expected_output_hint": "Should return Dan and Frank with their grades",
        "expected_rows": [("Dan", 85), ("Frank", 78)],
    },
    {
        "id": "easy_4",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE orders (id INTEGER, customer TEXT, amount REAL, status TEXT);",
        "setup": [
            "INSERT INTO orders VALUES (1,'Alice',250.0,'shipped');",
            "INSERT INTO orders VALUES (2,'Bob',89.0,'pending');",
            "INSERT INTO orders VALUES (3,'Carol',420.0,'shipped');",
        ],
        "broken_query": "SELECT customer, amount FROM orders WHERE status = 'shipped'",
        "fixed_query":  "SELECT customer, amount FROM orders WHERE status = 'shipped';",
        "error_message": "",
        "expected_output_hint": "Should return Alice and Carol with shipped order amounts",
        "expected_rows": [("Alice", 250.0), ("Carol", 420.0)],
    },
    {
        "id": "easy_5",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE inventory (id INTEGER, item TEXT, qty INTEGER, warehouse TEXT);",
        "setup": [
            "INSERT INTO inventory VALUES (1,'Bolts',500,'A');",
            "INSERT INTO inventory VALUES (2,'Nuts',300,'B');",
            "INSERT INTO inventory VALUES (3,'Screws',750,'A');",
        ],
        "broken_query": "SELECT item, qty FROM inventory WHERE warehouse = 'A' ORDR BY qty;",
        "fixed_query":  "SELECT item, qty FROM inventory WHERE warehouse = 'A' ORDER BY qty;",
        "error_message": "ParseError: near 'ORDR': syntax error",
        "expected_output_hint": "Should return Bolts and Screws ordered by quantity ascending",
        "expected_rows": [("Bolts", 500), ("Screws", 750)],
    },
    {
        "id": "easy_6",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE users (id INTEGER, username TEXT, age INTEGER, city TEXT);",
        "setup": [
            "INSERT INTO users VALUES (1,'alice',28,'Delhi');",
            "INSERT INTO users VALUES (2,'bob',35,'Mumbai');",
            "INSERT INTO users VALUES (3,'carol',22,'Delhi');",
        ],
        "broken_query": "SELECT username, age FORM users WHERE city = 'Delhi';",
        "fixed_query":  "SELECT username, age FROM users WHERE city = 'Delhi';",
        "error_message": "ParseError: near 'FORM': syntax error",
        "expected_output_hint": "Should return alice and carol from Delhi",
        "expected_rows": [("alice", 28), ("carol", 22)],
    },
    {
        "id": "easy_7",
        "task_id": "syntax_fix",
        "db_schema": "CREATE TABLE sales (id INTEGER, rep TEXT, amount REAL, region TEXT);",
        "setup": [
            "INSERT INTO sales VALUES (1,'Tom',15000,'North');",
            "INSERT INTO sales VALUES (2,'Sue',22000,'South');",
            "INSERT INTO sales VALUES (3,'Ray',18000,'North');",
        ],
        "broken_query": "SELECT rep, SUM(amount) FROM sales GRUP BY region;",
        "fixed_query":  "SELECT rep, SUM(amount) FROM sales GROUP BY region;",
        "error_message": "ParseError: near 'GRUP': syntax error",
        "expected_output_hint": "Should return total sales per region",
        "expected_rows": [("Tom", 33000.0), ("Sue", 22000.0)],
    },

    # ── MEDIUM: logic bugs ───────────────────────────────────────────
    {
        "id": "medium_1",
        "task_id": "logic_bug",
        "db_schema": "CREATE TABLE employees (id INTEGER, name TEXT, salary REAL, dept TEXT);",
        "setup": [
            "INSERT INTO employees VALUES (1,'Alice',75000,'Engineering');",
            "INSERT INTO employees VALUES (2,'Bob',50000,'Marketing');",
            "INSERT INTO employees VALUES (3,'Carol',90000,'Engineering');",
            "INSERT INTO employees VALUES (4,'Dave',45000,'Marketing');",
        ],
        "broken_query": "SELECT dept, AVG(salary) FROM employees GROUP BY name;",
        "fixed_query":  "SELECT dept, AVG(salary) FROM employees GROUP BY dept;",
        "error_message": "",
        "expected_output_hint": "Should return average salary per department, not per person",
        "expected_rows": [("Engineering", 82500.0), ("Marketing", 47500.0)],
    },
    {
        "id": "medium_2",
        "task_id": "logic_bug",
        "db_schema": "CREATE TABLE orders (id INTEGER, customer TEXT, amount REAL, status TEXT);",
        "setup": [
            "INSERT INTO orders VALUES (1,'Alice',250.0,'shipped');",
            "INSERT INTO orders VALUES (2,'Bob',89.0,'pending');",
            "INSERT INTO orders VALUES (3,'Carol',420.0,'shipped');",
            "INSERT INTO orders VALUES (4,'Dave',150.0,'cancelled');",
        ],
        "broken_query": "SELECT customer, amount FROM orders WHERE status != 'shipped';",
        "fixed_query":  "SELECT customer, amount FROM orders WHERE status = 'shipped';",
        "error_message": "",
        "expected_output_hint": "Should return only shipped orders — Alice and Carol",
        "expected_rows": [("Alice", 250.0), ("Carol", 420.0)],
    },
    {
        "id": "medium_3",
        "task_id": "logic_bug",
        "db_schema": "CREATE TABLE products (id INTEGER, name TEXT, price REAL, category TEXT);",
        "setup": [
            "INSERT INTO products VALUES (1,'Laptop',999.99,'Electronics');",
            "INSERT INTO products VALUES (2,'Shirt',29.99,'Clothing');",
            "INSERT INTO products VALUES (3,'Phone',699.99,'Electronics');",
            "INSERT INTO products VALUES (4,'Jeans',59.99,'Clothing');",
        ],
        "broken_query": "SELECT name, price FROM products WHERE price > 100 AND category = 'Clothing';",
        "fixed_query":  "SELECT name, price FROM products WHERE price > 100 AND category = 'Electronics';",
        "error_message": "",
        "expected_output_hint": "Should return expensive electronics — Laptop and Phone",
        "expected_rows": [("Laptop", 999.99), ("Phone", 699.99)],
    },
    {
        "id": "medium_4",
        "task_id": "logic_bug",
        "db_schema": "CREATE TABLE students (id INTEGER, name TEXT, score INTEGER, passed INTEGER);",
        "setup": [
            "INSERT INTO students VALUES (1,'Alice',85,1);",
            "INSERT INTO students VALUES (2,'Bob',45,0);",
            "INSERT INTO students VALUES (3,'Carol',72,1);",
            "INSERT INTO students VALUES (4,'Dave',38,0);",
        ],
        "broken_query": "SELECT COUNT(*) FROM students WHERE passed = 1 LIMIT 1;",
        "fixed_query":  "SELECT COUNT(*) FROM students WHERE passed = 1;",
        "error_message": "",
        "expected_output_hint": "Should return total count of passed students which is 2",
        "expected_rows": [(2,)],
    },
    {
        "id": "medium_5",
        "task_id": "logic_bug",
        "db_schema": "CREATE TABLE employees (id INTEGER, name TEXT, salary REAL, dept TEXT);",
        "setup": [
            "INSERT INTO employees VALUES (1,'Alice',75000,'Engineering');",
            "INSERT INTO employees VALUES (2,'Bob',50000,'Marketing');",
            "INSERT INTO employees VALUES (3,'Carol',90000,'Engineering');",
        ],
        "broken_query": "SELECT name, salary FROM employees ORDER BY salary ASC LIMIT 1;",
        "fixed_query":  "SELECT name, salary FROM employees ORDER BY salary DESC LIMIT 1;",
        "error_message": "",
        "expected_output_hint": "Should return the highest paid employee — Carol with 90000",
        "expected_rows": [("Carol", 90000.0)],
    },
    {
        "id": "medium_6",
        "task_id": "logic_bug",
        "db_schema": "CREATE TABLE sales (id INTEGER, rep TEXT, amount REAL, month INTEGER);",
        "setup": [
            "INSERT INTO sales VALUES (1,'Tom',15000,1);",
            "INSERT INTO sales VALUES (2,'Sue',22000,1);",
            "INSERT INTO sales VALUES (3,'Tom',18000,2);",
            "INSERT INTO sales VALUES (4,'Sue',25000,2);",
        ],
        "broken_query": "SELECT rep, amount FROM sales WHERE month = 1;",
        "fixed_query":  "SELECT rep, SUM(amount) FROM sales GROUP BY rep;",
        "error_message": "",
        "expected_output_hint": "Should return total sales per rep across all months",
        "expected_rows": [("Tom", 33000.0), ("Sue", 47000.0)],
    },

    # ── HARD: multi-table optimization ──────────────────────────────
    {
        "id": "hard_1",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE employees (id INTEGER, name TEXT, dept_id INTEGER, salary REAL);"
            "CREATE TABLE departments (id INTEGER, dept_name TEXT, budget REAL);"
        ),
        "setup": [
            "INSERT INTO departments VALUES (1,'Engineering',500000);",
            "INSERT INTO departments VALUES (2,'Marketing',200000);",
            "INSERT INTO employees VALUES (1,'Alice',1,75000);",
            "INSERT INTO employees VALUES (2,'Bob',2,50000);",
            "INSERT INTO employees VALUES (3,'Carol',1,90000);",
        ],
        "broken_query": "SELECT e.name, d.dept_name FROM employees e, departments d WHERE e.salary > 60000;",
        "fixed_query":  "SELECT e.name, d.dept_name FROM employees e JOIN departments d ON e.dept_id = d.id WHERE e.salary > 60000;",
        "error_message": "",
        "expected_output_hint": "Should return Alice and Carol with their department names using proper JOIN",
        "expected_rows": [("Alice", "Engineering"), ("Carol", "Engineering")],
    },
    {
        "id": "hard_2",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE orders (id INTEGER, customer_id INTEGER, amount REAL);"
            "CREATE TABLE customers (id INTEGER, name TEXT, city TEXT);"
        ),
        "setup": [
            "INSERT INTO customers VALUES (1,'Alice','Delhi');",
            "INSERT INTO customers VALUES (2,'Bob','Mumbai');",
            "INSERT INTO orders VALUES (1,1,250.0);",
            "INSERT INTO orders VALUES (2,1,180.0);",
            "INSERT INTO orders VALUES (3,2,420.0);",
        ],
        "broken_query": "SELECT c.name, o.amount FROM customers c LEFT JOIN orders o ON c.id = o.id;",
        "fixed_query":  "SELECT c.name, SUM(o.amount) FROM customers c JOIN orders o ON c.id = o.customer_id GROUP BY c.name;",
        "error_message": "",
        "expected_output_hint": "Should return total order amount per customer — Alice 430, Bob 420",
        "expected_rows": [("Alice", 430.0), ("Bob", 420.0)],
    },
    {
        "id": "hard_3",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE employees (id INTEGER, name TEXT, dept_id INTEGER, salary REAL);"
            "CREATE TABLE departments (id INTEGER, dept_name TEXT, budget REAL);"
        ),
        "setup": [
            "INSERT INTO departments VALUES (1,'Engineering',500000);",
            "INSERT INTO departments VALUES (2,'Marketing',200000);",
            "INSERT INTO employees VALUES (1,'Alice',1,75000);",
            "INSERT INTO employees VALUES (2,'Bob',2,50000);",
            "INSERT INTO employees VALUES (3,'Carol',1,90000);",
            "INSERT INTO employees VALUES (4,'Dave',2,45000);",
        ],
        "broken_query": "SELECT dept_name, COUNT(*) FROM departments GROUP BY dept_name;",
        "fixed_query":  "SELECT d.dept_name, COUNT(e.id) FROM departments d JOIN employees e ON d.id = e.dept_id GROUP BY d.dept_name;",
        "error_message": "",
        "expected_output_hint": "Should return headcount per department — Engineering 2, Marketing 2",
        "expected_rows": [("Engineering", 2), ("Marketing", 2)],
    },
    {
        "id": "hard_4",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE products (id INTEGER, name TEXT, category_id INTEGER, price REAL);"
            "CREATE TABLE categories (id INTEGER, cat_name TEXT);"
            "CREATE TABLE order_items (id INTEGER, product_id INTEGER, qty INTEGER);"
        ),
        "setup": [
            "INSERT INTO categories VALUES (1,'Electronics');",
            "INSERT INTO categories VALUES (2,'Clothing');",
            "INSERT INTO products VALUES (1,'Laptop',1,999.99);",
            "INSERT INTO products VALUES (2,'Shirt',2,29.99);",
            "INSERT INTO products VALUES (3,'Phone',1,699.99);",
            "INSERT INTO order_items VALUES (1,1,2);",
            "INSERT INTO order_items VALUES (2,3,5);",
            "INSERT INTO order_items VALUES (3,2,10);",
        ],
        "broken_query": "SELECT p.name, oi.qty FROM products p JOIN order_items oi ON p.id = oi.id;",
        "fixed_query":  "SELECT p.name, SUM(oi.qty) as total_qty FROM products p JOIN order_items oi ON p.id = oi.product_id GROUP BY p.name ORDER BY total_qty DESC;",
        "error_message": "",
        "expected_output_hint": "Should return total quantity ordered per product, highest first",
        "expected_rows": [("Shirt", 10), ("Phone", 5), ("Laptop", 2)],
    },
    {
        "id": "hard_5",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE employees (id INTEGER, name TEXT, manager_id INTEGER, salary REAL);"
        ),
        "setup": [
            "INSERT INTO employees VALUES (1,'CEO',NULL,200000);",
            "INSERT INTO employees VALUES (2,'Alice',1,90000);",
            "INSERT INTO employees VALUES (3,'Bob',1,85000);",
            "INSERT INTO employees VALUES (4,'Carol',2,70000);",
        ],
        "broken_query": "SELECT e.name, m.name FROM employees e JOIN employees m ON e.id = m.manager_id;",
        "fixed_query":  "SELECT e.name, m.name as manager FROM employees e JOIN employees m ON e.manager_id = m.id WHERE e.manager_id IS NOT NULL;",
        "error_message": "",
        "expected_output_hint": "Should return each employee with their manager name — self join",
        "expected_rows": [("Alice", "CEO"), ("Bob", "CEO"), ("Carol", "Alice")],
    },
    {
        "id": "hard_6",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE orders (id INTEGER, customer_id INTEGER, amount REAL, status TEXT);"
            "CREATE TABLE customers (id INTEGER, name TEXT, tier TEXT);"
        ),
        "setup": [
            "INSERT INTO customers VALUES (1,'Alice','gold');",
            "INSERT INTO customers VALUES (2,'Bob','silver');",
            "INSERT INTO customers VALUES (3,'Carol','gold');",
            "INSERT INTO orders VALUES (1,1,500,'shipped');",
            "INSERT INTO orders VALUES (2,2,200,'shipped');",
            "INSERT INTO orders VALUES (3,3,800,'shipped');",
            "INSERT INTO orders VALUES (4,1,300,'pending');",
        ],
        "broken_query": "SELECT c.name, SUM(o.amount) FROM customers c JOIN orders o ON c.id = o.customer_id GROUP BY c.name;",
        "fixed_query":  "SELECT c.name, SUM(o.amount) FROM customers c JOIN orders o ON c.id = o.customer_id WHERE c.tier = 'gold' AND o.status = 'shipped' GROUP BY c.name;",
        "error_message": "",
        "expected_output_hint": "Should return total shipped order amounts for gold tier customers only",
        "expected_rows": [("Alice", 500.0), ("Carol", 800.0)],
    },
    {
        "id": "hard_7",
        "task_id": "multi_table",
        "db_schema": (
            "CREATE TABLE employees (id INTEGER, name TEXT, dept_id INTEGER, salary REAL);"
            "CREATE TABLE departments (id INTEGER, dept_name TEXT, budget REAL);"
        ),
        "setup": [
            "INSERT INTO departments VALUES (1,'Engineering',500000);",
            "INSERT INTO departments VALUES (2,'Marketing',200000);",
            "INSERT INTO employees VALUES (1,'Alice',1,75000);",
            "INSERT INTO employees VALUES (2,'Bob',2,50000);",
            "INSERT INTO employees VALUES (3,'Carol',1,90000);",
            "INSERT INTO employees VALUES (4,'Dave',2,45000);",
        ],
        "broken_query": "SELECT dept_name FROM departments WHERE budget > AVG(budget);",
        "fixed_query":  "SELECT dept_name FROM departments WHERE budget > (SELECT AVG(budget) FROM departments);",
        "error_message": "misuse of aggregate function AVG()",
        "expected_output_hint": "Should return departments with above-average budget — Engineering only",
        "expected_rows": [("Engineering",)],
    },
]
# module-level globals — persist across all instances
_CURRENT_SCENARIO = None
_CURRENT_STEP = 0

def compute_f1(predicted_rows, expected_rows):
    if not expected_rows:
        return 1.0 if not predicted_rows else 0.0
    pred_set = [tuple(str(v) for v in row) for row in predicted_rows]
    exp_set  = [tuple(str(v) for v in row) for row in expected_rows]
    pred_multiset = {}
    for row in pred_set:
        pred_multiset[row] = pred_multiset.get(row, 0) + 1
    exp_multiset = {}
    for row in exp_set:
        exp_multiset[row] = exp_multiset.get(row, 0) + 1
    true_positives = 0
    for row, count in exp_multiset.items():
        true_positives += min(count, pred_multiset.get(row, 0))
    precision = true_positives / len(pred_set) if pred_set else 0.0
    recall    = true_positives / len(exp_set)  if exp_set  else 0.0
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


def run_query_safe(db_schema, setup_stmts, query):
    try:
        conn = sqlite3.connect(":memory:")
        cur  = conn.cursor()
        for stmt in db_schema.split(";"):
            stmt = stmt.strip()
            if stmt:
                cur.execute(stmt)
        for stmt in setup_stmts:
            cur.execute(stmt)
        conn.commit()
        cur.execute(query)
        rows = cur.fetchall()
        conn.close()
        return rows, ""
    except Exception as e:
        return [], str(e)


# module-level globals — survive across all instances
_CURRENT_SCENARIO = random.choice(SCENARIOS)
_CURRENT_STEP     = 0


class SqlQueryDebuggerEnvironment(Environment):

    SUPPORTS_CONCURRENT_SESSIONS: bool = False
    MAX_STEPS = 5

    def __init__(self):
        self._state = State(episode_id=str(uuid4()), step_count=0)

    def reset(self, task_id: str = None) -> SqlQueryDebuggerObservation:
        global _CURRENT_SCENARIO, _CURRENT_STEP
        self._state = State(episode_id=str(uuid4()), step_count=0)
        _CURRENT_STEP = 0

        pool = [s for s in SCENARIOS if s["task_id"] == task_id] if task_id else SCENARIOS
        _CURRENT_SCENARIO = random.choice(pool)

        sample_rows, _ = run_query_safe(
            _CURRENT_SCENARIO["db_schema"],
            _CURRENT_SCENARIO["setup"],
            "SELECT * FROM " + _CURRENT_SCENARIO["db_schema"]
                .split("CREATE TABLE ")[1].split(" ")[0] + " LIMIT 3;"
        )

        return SqlQueryDebuggerObservation(
            broken_query         = _CURRENT_SCENARIO["broken_query"],
            db_schema               = _CURRENT_SCENARIO["db_schema"],
            error_message        = _CURRENT_SCENARIO["error_message"],
            sample_rows          = json.dumps(sample_rows),
            expected_output_hint = _CURRENT_SCENARIO["expected_output_hint"],
            task_id              = _CURRENT_SCENARIO["task_id"],
            attempts_remaining   = self.MAX_STEPS,
            last_result          = None,
            done                 = False,
            reward               = 0.0,
        )

    def step(self, action: SqlQueryDebuggerAction) -> SqlQueryDebuggerObservation:
        global _CURRENT_SCENARIO, _CURRENT_STEP
        _CURRENT_STEP += 1
        self._state.step_count = _CURRENT_STEP
        attempts_left = self.MAX_STEPS - _CURRENT_STEP

        rows, error = run_query_safe(
            _CURRENT_SCENARIO["db_schema"],
            _CURRENT_SCENARIO["setup"],
            action.fixed_query,
        )

        f1     = compute_f1(rows, _CURRENT_SCENARIO["expected_rows"])
        done   = f1 >= 0.99 or attempts_left <= 0
        bonus  = 0.1 * (attempts_left / self.MAX_STEPS) if f1 >= 0.99 else 0.0
        reward = min(round(f1 + bonus, 4), 1.0)

        return SqlQueryDebuggerObservation(
            broken_query         = _CURRENT_SCENARIO["broken_query"],
            db_schema               = _CURRENT_SCENARIO["db_schema"],
            error_message        = error,
            sample_rows          = json.dumps(_CURRENT_SCENARIO["setup"]),
            expected_output_hint = _CURRENT_SCENARIO["expected_output_hint"],
            task_id              = _CURRENT_SCENARIO["task_id"],
            attempts_remaining   = max(attempts_left, 0),
            last_result          = json.dumps(rows),
            done                 = done,
            reward               = reward,
        )

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