sql-debug-env / README.md
aakashdandekar's picture
Add Hugging Face Space configuration to README
ae209be
|
Raw History Blame Contribute Delete
14.7 kB
---
title: sql-debug-env
emoji: 🐳
colorFrom: blue
colorTo: indigo
sdk: docker
app_port: 7860
---
# SQL Debug Environment
**An OpenEnv-compliant real-world environment where AI agents learn to debug and fix broken SQL queries.**
[![OpenEnv](https://img.shields.io/badge/OpenEnv-1.0.0-blue)](https://openenv.dev)
[![Python 3.11](https://img.shields.io/badge/python-3.11-blue.svg)](https://python.org)
[![FastAPI](https://img.shields.io/badge/FastAPI-0.115-green)](https://fastapi.tiangolo.com)
[![Docker](https://img.shields.io/badge/docker-ready-blue)](https://docker.com)
---
## 🧠 Motivation
SQL debugging is one of the most common, high-value tasks performed by developers, data engineers, and analysts every day. A study by Stack Overflow found SQL is consistently among the top 5 most-used languages. Yet existing RL/agent benchmarks almost entirely focus on SQL *generation*, not *repair*.
This environment fills that gap. Given a broken SQL query, an agent must:
1. Analyze the schema, sample data, and error message
2. Identify what's wrong (syntax error, wrong column, bad JOIN, logic bug, etc.)
3. Submit a corrected query that runs without error AND returns the expected results
This is a genuine real-world task. Every improvement in SQL debugging agents has immediate practical value.
---
## πŸ—‚οΈ Project Structure
```
sql-query-env/
β”œβ”€β”€ openenv.yaml # OpenEnv environment manifest
β”œβ”€β”€ models.py # Typed Pydantic models: Action, Observation, State
β”œβ”€β”€ tasks.py # All task definitions (9 tasks across 3 difficulty levels)
β”œβ”€β”€ graders.py # Deterministic graders scoring 0.0–1.0
β”œβ”€β”€ inference.py # Baseline inference script (OpenAI client)
β”œβ”€β”€ Dockerfile # Container for HF Spaces deployment
β”œβ”€β”€ pyproject.toml # Package metadata
β”œβ”€β”€ server/
β”‚ β”œβ”€β”€ app.py # FastAPI HTTP server (step/reset/state endpoints)
β”‚ β”œβ”€β”€ sql_environment.py# Core environment class
β”‚ └── requirements.txt # Python dependencies
└── README.md
```
---
## 🌍 Environment Description
The **SQL Debug Environment** presents an agent with a broken SQL query and all context needed to fix it. The agent iteratively submits corrected SQL queries and receives graded rewards based on correctness.
### Real-World Task
SQL query debugging:
- **Syntax errors**: Missing keywords (`FROM`, `GROUP BY`), wrong column names
- **Semantic/logic errors**: Wrong JOIN conditions, incorrect aggregation functions
- **Complex logic bugs**: Wrong window function offsets, missing `LEFT JOIN`, broken CTEs
---
## βš™οΈ How It Works
This project implements a Reinforcement Learning (RL) loop adapted for an LLM agent interacting over an API. The environment is served by a **FastAPI backend** (`server/app.py`).
Here is the step-by-step flow of a single "episode" (a debugging session):
1. **Reset (Start an Episode)**: The agent sends a `POST /reset` request to the server, picking a difficulty level (easy, medium, hard).
2. **Observation (The Context)**: The server returns the initial state containing everything the agent needs:
- **`buggy_query`**: The broken SQL string.
- **`error_message`**: The error thrown by the database when running the buggy query (or empty if it's a silent semantic bug).
- **`schema_sql`**: The `CREATE TABLE` definitions.
- **`sample_data`**: Small `INSERT` statements to show the data shape.
- **`task_description`**: A natural language instruction of what the query is *supposed* to do.
3. **Step (Take an Action)**: The agent generates what it believes is the fixed query and submits it via `POST /step`.
4. **Grading & Reward**: The backend (`graders.py` and `sql_environment.py`) sets up an **in-memory SQLite database**, loads the schema and sample data, and executes the agent's SQL query.
- It checks for execution errors.
- It checks if the returned columns match the expected columns.
- It checks if the returned data tuples match the expected golden data exactly.
5. **Feedback Loop**: Depending on the results, the environment returns a **shaped reward** (e.g., `+1.0` if perfect, `+0.5` if it ran but data is wrong, `+0.1` if it at least fixed the syntax error) and an updated observation.
6. The agent gets up to **5 attempts** to achieve a 1.0 success score.
### Key Technical Components
- **The Grader System (`graders.py`)**: This engine guarantees *deterministic evaluation*. Because it uses SQLite in memory, it runs entirely isolated for every attempt. It compares the output of the agent's query against a hardcoded set of `expected_rows`.
- **Reward Shaping**: Instead of a simple pass/fail, the environment gives *partial* points. If the agent fixes a syntax error but gets the logical result wrong, it still earns a small reward. There is also a time-decaying bonus for solving the bug on earlier attempts.
- **Baseline Inference (`inference.py`)**: A provided Python script that demonstrates how an LLM acts as the agent, parsing the observation spaces, planning fixes, constructing prompts, and making API calls to the server.
---
## πŸ“ Action Space
```python
class Action(BaseModel):
query: str # The corrected SQL query to execute
explanation: str # Optional explanation of what was fixed
```
The agent submits a corrected SQL string. The environment executes it against an in-memory SQLite database and evaluates correctness.
---
## πŸ‘οΈ Observation Space
```python
class Observation(BaseModel):
schema_sql: str # CREATE TABLE statements for all tables
sample_data: str # INSERT statements showing sample rows
buggy_query: str # The original broken SQL query
error_message: str # Error from running the buggy query (empty = semantic bug)
task_description: str # Natural language description of what the query should do
expected_columns: list # Column names expected in the correct result
attempts_remaining: int # How many fix attempts remain (max 5)
last_attempt_result: str # Feedback from last submitted query
reward: float # Reward from last action
done: bool # Whether the episode is complete
success: bool # Whether the task was solved correctly
```
---
## 🧩 Tasks
### 3 Task Groups with 9 Total Tasks
| Task Group | Difficulty | Description | Expected Baseline Score |
|---|---|---|---|
| `easy_syntax_fix` | Easy | Fix syntax errors: missing FROM, wrong column names, missing GROUP BY | ~0.75 |
| `medium_logic_fix` | Medium | Fix semantic errors: wrong JOINs, wrong aggregates, INNER vs LEFT JOIN | ~0.50 |
| `hard_optimization_fix` | Hard | Fix complex queries with CTEs, window functions, correlated subqueries | ~0.25 |
### Easy Tasks
- **`easy_missing_from`** β€” Missing `FROM` keyword in a simple SELECT query
- **`easy_wrong_column`** β€” Query uses `wage` instead of `salary` column name
- **`easy_missing_groupby`** β€” `HAVING` without `GROUP BY` causes aggregate error
### Medium Tasks
- **`medium_wrong_join`** β€” JOIN on `e.id = d.id` should be `e.dept_id = d.id`
- **`medium_wrong_aggregate`** β€” Uses `SUM` instead of `AVG` for average calculation
- **`medium_inner_vs_left`** β€” `INNER JOIN` excludes customers with no orders; needs `LEFT JOIN + COALESCE`
### Hard Tasks
- **`hard_window_function`** β€” `ROW_NUMBER() OVER (... ORDER BY salary DESC)` needs tiebreaker `id ASC`
- **`hard_correlated_subquery`** β€” Incorrectly correlated subquery in `NOT IN` clause
- **`hard_multi_cte`** β€” Multi-CTE with `LAG(total_revenue, 2)` should be `LAG(total_revenue, 1)`
---
## 🎯 Reward Function
The reward function provides **meaningful partial progress signals** throughout each episode β€” not just binary end-of-episode rewards.
```
Reward per step:
+1.0–1.2 Query runs AND returns exactly expected rows (early-solve bonus: +0.05 per unused attempt)
+0.50 Query runs AND has correct columns but wrong data
+0.20 Query runs AND has correct number of rows (wrong data)
+0.10 Query runs but has wrong columns (at least syntactically valid)
0.00 Query fails with same/new error (no penalty on first attempt)
-0.05 Wasted attempt: query still fails after attempt 1
Episode ends when: score=1.0 (success) OR max_steps=5 reached (failure)
```
**Why this reward shape works for RL:**
- Agents get signal even on partial fixes (fixed the error, wrong logic β†’ +0.1)
- Early solving is rewarded (faster fix = higher cumulative reward)
- Discourages random guessing (repeated failures cost -0.05)
- Dense signal across the full 5-step trajectory
---
## πŸš€ Quick Start
### Using Docker
```bash
# Build image
docker build -t sql-query-env .
# Run container
docker run -p 8000:8000 sql-query-env
# Test
curl http://localhost:8000/health
```
### Local Setup (Python)
```bash
# Create Virtual Environment
python3 -m venv .venv && source .venv/bin/activate
# Install
pip install -r server/requirements.txt
# Run server
python -m uvicorn server.app:app --host 0.0.0.0 --port 8000
# Test
curl http://localhost:8000/health
```
---
## πŸ”Œ API Usage
### Reset an Episode
```bash
curl -X POST http://localhost:8000/reset \
-H "Content-Type: application/json" \
-d '{"task_group": "easy_syntax_fix", "seed": 42}'
```
Response includes `session_id`, initial `observation`, and `state`.
### Take a Step
```bash
curl -X POST http://localhost:8000/step \
-H "Content-Type: application/json" \
-d '{
"session_id": "<from reset>",
"query": "SELECT first_name, last_name FROM employees WHERE department_id = 10 ORDER BY last_name;",
"explanation": "Added missing FROM keyword"
}'
```
### Get State
```bash
curl "http://localhost:8000/state?session_id=<session_id>"
```
### Grade a Query Directly
```bash
curl -X POST http://localhost:8000/grade \
-H "Content-Type: application/json" \
-d '{"task_id": "easy_missing_from", "query": "SELECT first_name, last_name FROM employees WHERE department_id = 10 ORDER BY last_name;"}'
```
---
## πŸ€– Baseline Inference
```bash
# Set environment variables
export API_BASE_URL="https://api.openai.com/v1"
export MODEL_NAME="gpt-4o-mini"
export HF_TOKEN="your-api-key"
export ENV_BASE_URL="http://localhost:8000"
# Run baseline (requires environment to be running)
python inference.py
```
### Baseline Scores (GPT-4o-mini)
| Task Group | Avg Score | Success Rate |
|---|---|---|
| `easy_syntax_fix` | 0.78 | 78% |
| `medium_logic_fix` | 0.51 | 51% |
| `hard_optimization_fix` | 0.24 | 24% |
| **Overall** | **0.51** | **51%** |
---
## πŸ“š Step-by-Step Example Explanations
The project formerly included example scripts to demonstrate usage. Their step-by-step methodologies are outlined below for reference:
### Step 1: Programmatic Interaction (Python Client)
To programmatically interact with the environment, developers can use the provided `client.py`. This involves initializing an asynchronous client connected to the local server, resetting the environment to fetch a target SQL task (and the buggy query), computing a fix, and submitting an action back to the server. The server then responds with shaped rewards and a success boolean based on the correctness of the generated SQL.
### Step 2: Baseline LLM Evaluation
To evaluate an LLM's baseline coding skills, users can write a script that authenticates with an OpenAI-compatible API endpoint. The script iterates over the defined SQL task groups (`easy`, `medium`, `hard`), dynamically constructs a prompt containing the database schema and sample data, requests a SQL patch from the LLM, and submits it to the environment to calculate the total reward trajectory and success rate.
### Step 3: Reinforcement Learning Training (GRPO)
For training an AI agent, the environment seamlessly integrates with Hugging Face's TRL (Transformer Reinforcement Learning) library. Using the OpenEnv abstraction, users can set up a Group Relative Policy Optimization (GRPO) training loop. The training script connects to the local environment and uses the partial and full rewards generated by the environment's deterministic grader to actively shape the policy of a base instruct-model, teaching it to iteratively write syntactically and logically correct SQL.
---
## ☁️ Hugging Face Spaces Deployment
This environment is deployed as a Docker Space on Hugging Face.
**Space URL**: `https://huggingface.co/spaces/<your-username>/sql-debug-env`
The space:
- Responds to HTTP requests on port 7860 (HF default)
- Tagged with `openenv` for discoverability
- Returns 200 on `GET /health`
- Responds correctly to `POST /reset`
To deploy your own:
```bash
# Login
huggingface-cli login
# Create space
huggingface-cli repo create sql-debug-env --type space --space_sdk docker
# Push
git remote add hf https://huggingface.co/spaces/<username>/sql-debug-env
git push hf main
```
---
## βœ… Pre-Submission Checklist
- [x] HF Space deploys and returns 200 on health ping
- [x] OpenEnv spec: `openenv.yaml`, typed models, `step()`/`reset()`/`state()` endpoints
- [x] Dockerfile builds and runs cleanly (`docker build && docker run`)
- [x] `inference.py` in root, uses OpenAI client, reads `API_BASE_URL`/`MODEL_NAME`/`HF_TOKEN`
- [x] 3+ task groups with graders scoring 0.0–1.0
- [x] Graders are deterministic and reproducible
- [x] Reward function provides partial progress signals (not just sparse)
- [x] Runtime < 20 min on vcpu=2, 8gb RAM
- [x] README includes environment description, action/observation spaces, setup instructions
---
## πŸ§ͺ Validation
Run the local validation script before submitting:
```bash
python validate.py
```
This checks:
- All 9 graders return 1.0 on correct solutions
- All 9 graders return 0.0 on completely wrong queries
- Environment reset/step/state cycle works correctly
- All expected_rows match actual SQLite output
---
## πŸ“Š Why This Environment Is Novel
Existing OpenEnv environments focus on: email triage, customer support, scheduling. **No existing environment targets SQL repair** β€” a daily pain point for millions of engineers.
Key novelty:
- Uses **SQLite in-memory databases** for fully deterministic, reproducible grading
- **9 hand-crafted bug scenarios** spanning syntax β†’ logic β†’ complex query levels
- Graders check **both** execution success AND result-set correctness
- Reward shaping designed specifically for **multi-attempt iterative refinement**
- Tasks are drawn from **real bugs** commonly found in production SQL codebases
---
## πŸ“„ License
Licensed under the **Apache License 2.0**. See [LICENSE](LICENSE) for full terms.