deepakvish001/sql-repair-env

SQL Query Repair Environment

An OpenEnv-compliant reinforcement learning environment where AI agents learn to identify and fix broken SQL queries — one of the most common, real-world tasks performed by developers and data engineers every day.

Why This Environment?

Every developer who has worked with databases has encountered broken SQL queries. They arise from:

• —Typos (missing commas, wrong keywords)
• —Logic errors (wrong JOIN type, wrong column in ON clause)
• —Aggregation mistakes (bad GROUP BY, wrong computation)

Training agents to repair SQL teaches systematic debugging — a transferable skill across many domains. There is currently no standard benchmark for this capability.

Environment Overview

Observation Space

Each observation contains:

Action Space

Reward Space

Tasks

Task 1 — Fix Syntax Errors (Easy)

Schema: employees(id, name, department, salary, hire_date)

Broken query:

SELECT name department salary
FROM employees
WHERE salary > 60000
ORDER BY salary DESC

Problem: Missing commas between column names — a common beginner mistake.

Grader criteria (0.0 → 1.0):

• —Syntax valid (executes): +0.30
• —Correct columns selected: +0.30
• —Correct row count (7 employees with salary > 60k): +0.25
• —Correct ORDER BY (salary DESC): +0.15

Expected baseline score: ~0.85

Task 2 — Fix Logic Errors (Medium)

Schema: customers(id, name, email, city) + orders(id, customer_id, total_amount, order_date, status)

Broken query:

SELECT c.name, c.email, COUNT(o.id) AS order_count
FROM customers c
INNER JOIN orders o ON c.id = o.product_id   -- ← two bugs
GROUP BY c.id, c.name, c.email
ORDER BY order_count DESC

Problems:

1. 1.INNER JOIN should be LEFT JOIN (loses customers with zero orders)
2. 2.o.product_id should be o.customer_id in the JOIN condition

Grader criteria:

• —Executes: +0.20
• —Returns all 5 customers (LEFT JOIN): +0.35
• —Customers with 0 orders show count=0: +0.30
• —Correct columns: +0.15

Expected baseline score: ~0.55

Task 3 — Fix Complex Aggregation Query (Hard)

Schema: products(id, name, category, price, stock) + sales(id, product_id, quantity, sale_date, region)

Broken query:

SELECT p.category, p.name, SUM(s.quantity) AS revenue
FROM products p
INNER JOIN sales s ON p.id = s.product_id
GROUP BY p.category            -- ← Bug 1
ORDER BY revenue ASC           -- ← Bug 3

Three bugs:

1. 1.GROUP BY p.category → should be GROUP BY p.category, p.name
2. 2.SUM(s.quantity) → should be SUM(s.quantity * p.price) (revenue, not units)
3. 3.ORDER BY revenue ASC → should be ORDER BY revenue DESC

Grader criteria:

• —Executes: +0.20
• —5 rows returned (one per product): +0.20
• —Revenue uses price multiplication (values > 1000): +0.25
• —Correct descending order: +0.20
• —Top row is Laptop Pro (~$10,399): +0.15

Expected baseline score: ~0.40

HTTP API

Setup & Usage

Local (Python)

git clone <your-repo>
cd sql-repair-env

pip install -r requirements.txt

# Start the server
uvicorn server:app --host 0.0.0.0 --port 7860

# Run baseline agent
export HF_TOKEN=your_hf_token
export MODEL_NAME=meta-llama/Llama-3.3-70B-Instruct
export API_BASE_URL=https://router.huggingface.co/v1
python inference.py

Docker

docker build -t sql-repair-env .
docker run -p 7860:7860 \
  -e HF_TOKEN=$HF_TOKEN \
  -e MODEL_NAME=$MODEL_NAME \
  -e API_BASE_URL=$API_BASE_URL \
  sql-repair-env

Quick API Test

# Reset to easy task
curl -X POST http://localhost:7860/reset \
  -H "Content-Type: application/json" \
  -d '{"task_id": "fix_syntax_errors"}'

# Submit a fixed query
curl -X POST http://localhost:7860/step \
  -H "Content-Type: application/json" \
  -d '{"fixed_query": "SELECT name, department, salary FROM employees WHERE salary > 60000 ORDER BY salary DESC"}'

Baseline Scores

Measured with meta-llama/Llama-3.3-70B-Instruct via HF Inference Router:

Reward Design Notes

The reward function is designed to give continuous signal across the trajectory:

• —Agents that produce syntactically invalid queries get 0.0 — strong signal to fix syntax first
• —Agents that fix syntax but have wrong logic get 0.20–0.60 — gradient toward correct logic
• —Progressive hints are revealed after each failed attempt to guide the agent
• —Logic-correct answers end the episode early (efficiency bonus implicit in fewer steps used)
• —The hardest task requires fixing 3 independent bugs — partial credit is awarded per bug fixed

This design avoids the sparse reward problem and gives agents useful signal even on hard tasks.

Project Structure

sql-repair-env/
├── Dockerfile
├── README.md
├── openenv.yaml
├── requirements.txt
├── inference.py          ← baseline agent
├── server.py             ← FastAPI HTTP server
└── env/
    ├── __init__.py
    ├── models.py         ← Pydantic Observation / Action / Reward / State
    ├── database.py       ← SQLite setup + seed data
    ├── tasks.py          ← Task definitions + deterministic graders
    └── environment.py    ← Core environment (reset / step / state)