Spaces:

aakashdandekar2006
/

sql-debug-env

Sleeping

App Files Files Community

sql-debug-env / README.md

aakashdandekar

Add Hugging Face Space configuration to README

ae209be 7 months ago

preview code

Raw History Blame Contribute Delete

14.7 kB

metadata

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.

🧠 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:

Analyze the schema, sample data, and error message
Identify what's wrong (syntax error, wrong column, bad JOIN, logic bug, etc.)
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):

Reset (Start an Episode): The agent sends a POST /reset request to the server, picking a difficulty level (easy, medium, hard).
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.
Step (Take an Action): The agent generates what it believes is the fixed query and submits it via POST /step.
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.
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.
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

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

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

# 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)

# 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

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

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

curl "http://localhost:8000/state?session_id=<session_id>"

Grade a Query Directly

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

# 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:

# 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

HF Space deploys and returns 200 on health ping
OpenEnv spec: openenv.yaml, typed models, step()/reset()/state() endpoints
Dockerfile builds and runs cleanly (docker build && docker run)
inference.py in root, uses OpenAI client, reads API_BASE_URL/MODEL_NAME/HF_TOKEN
3+ task groups with graders scoring 0.0–1.0
Graders are deterministic and reproducible
Reward function provides partial progress signals (not just sparse)
Runtime < 20 min on vcpu=2, 8gb RAM
README includes environment description, action/observation spaces, setup instructions

🧪 Validation

Run the local validation script before submitting:

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 for full terms.