Apppublic

Kalletlamadhav/sql-optimization-env

sourceHugging Faceupdated 6mo agoView on Hugging Face

0likes

App README

🗄️ SQL Query Optimization — RL Training Environment

Meta PyTorch OpenEnv Hackathon × SST 2026 Author: Kalletla Madhav OpenEnv-compliant · SQLite · Indian Data Domains · Curriculum Learning

Overview

A containerised FastAPI service that exposes the OpenEnv standard API (/reset, /step, /state) for training LLM agents to optimise slow SQL queries.

The environment actually executes queries on a real SQLite database populated with synthetic Indian enterprise data (GST, PDS/Ration, IRCTC, MGNREGA) and grades agents on real wall-clock speedup — no simulations, no fake timing.

Quick Start

bash

# Clone and build
git clone https://huggingface.co/spaces/Kalletlamadhav/sql-optimization-env
cd sql-optimization-env

# Generate the benchmark database
python data/seed_database.py --rows 100000

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

# Run the baseline inference script
export OPENAI_API_KEY=sk-...
export MODEL_NAME=gpt-4o-mini
python inference.py

Or with Docker:

bash

docker build -t sql-opt-env .
docker run -p 7860:7860 -e OPENAI_API_KEY=sk-... sql-opt-env

API Reference

`POST /reset`

Start a new episode. Optionally pin to a specific task.

json

{ "task_id": "gst_missing_index" }

`POST /step`

Submit an optimised query.

json

{
  "optimized_query": "SELECT invoice_id, invoice_date FROM gst_invoice_records WHERE gstin_supplier = '27AABCU9603R1ZX' AND invoice_date >= '2025-01-01'",
  "identified_pattern": "MISSING_INDEX",
  "explanation": "The query does a full table scan because gstin_supplier has no index.",
  "index_statements": [
    "CREATE INDEX idx_gst_supplier_date ON gst_invoice_records(gstin_supplier, invoice_date DESC)"
  ],
  "schema_analysis": "Table has 100k rows; gstin_supplier is high-selectivity but unindexed."
}

`GET /state`

Returns full environment state including curriculum level and episode history.

Observation Space

Field	Type	Description
`task_id`	string	Current task identifier
`step_number`	int	Current step within episode
`goal`	string	Natural language optimisation objective
`schema_ddl`	string	CREATE TABLE statements for all involved tables
`current_query`	string	The slow SQL query to optimise
`execution_plan`	object	Structured EXPLAIN QUERY PLAN output
`execution_time_ms`	float	Wall-clock execution time in milliseconds
`row_count`	int	Rows returned by current query
`db_stats`	object	Per-table row counts and existing indexes
`curriculum_level`	int	Current difficulty level (1–5)
`anti_pattern_hint`	string	Hint provided at levels 1–2 only

Action Space

Field	Type	Description
`optimized_query`	string	The agent's rewritten SQL
`identified_pattern`	enum	One of 7 anti-pattern types
`explanation`	string	Agent's reasoning for the fix
`index_statements`	list[string]	Optional CREATE INDEX statements
`schema_analysis`	string	Agent's schema observations

Anti-pattern enum values: N_PLUS_ONE, CARTESIAN_PRODUCT, MISSING_INDEX, SELECT_STAR, LEADING_WILDCARD, IMPLICIT_CAST, UNBOUNDED_AGGREGATION, NONE

Reward Function

Dimension	Weight	Description
Speedup	0.35	`log10(orig_time / opt_time)` scaled to 0–0.35
Equivalence	0.25	Jaccard similarity of result sets
Pattern	0.20	Correct identification + explanation quality
Index	0.10	Index created AND used by query planner
Simplicity	0.10	Optimised query ≤ original length

Penalties: syntax error (−0.30), wrong results (−0.20), slower query (−0.10), timeout (−0.05/s), reward hacking (−0.40).

Tasks

ID	Level	Anti-Pattern	Domain
`pds_select_star`	1 — Intro	SELECT_STAR	PDS Ration
`gst_missing_index`	2 — Easy	MISSING_INDEX	GST
`railway_simple_filter`	2 — Easy	MISSING_INDEX	IRCTC
`gst_n_plus_one`	3 — Medium	NPLUSONE	GST
`pds_cartesian`	3 — Medium	CARTESIAN_PRODUCT	PDS Ration
`mgnrega_wildcard`	3 — Medium	LEADING_WILDCARD	MGNREGA
`gst_multi_join`	4 — Hard	NPLUSONE	GST
`railway_tatkal_workload`	4 — Hard	NPLUSONE + CARTESIAN	IRCTC
`mgnrega_schema_e`	5 — Expert	UNBOUNDED_AGGREGATION	MGNREGA

Curriculum Learning

The environment automatically adjusts difficulty:

—Graduate up after 3 consecutive episodes with score ≥ 0.70
—Step down after 2 consecutive episodes with score < 0.40
—State persisted to /tmp/curriculum_state.json across restarts

Baseline Scores

Scores from gpt-4o-mini running inference.py (SEED=42):

Task	Steps	Final Reward
`gst_missing_index`	1	TBD
`gst_n_plus_one`	2	TBD
`gst_multi_join`	3	TBD
Average	—	TBD

Run `python inference.py` to reproduce.

Data Domains

Schema	Domain
GST	Indian GST e-invoicing (100k invoices)
PDS	Public Distribution System ration cards
IRCTC	Indian Railway PNR bookings
MGNREGA	Rural employment attendance + payments

All data is synthetic, seeded with SEED=42 for full reproducibility.

Project Structure

sql-optimization-env/
├── openenv.yaml
├── Dockerfile
├── inference.py
├── requirements.txt
├── server/
│   ├── main.py
│   ├── environment.py
│   ├── models.py
│   ├── reward.py
│   ├── hack_detector.py
│   └── graders/
├── tasks/
│   ├── easy/   (3 tasks)
│   ├── medium/ (3 tasks)
│   └── hard/   (3 tasks)
├── data/
│   ├── schemas/
│   └── fixtures/benchmark_seed42.db
└── curriculum/

Validation

bash

pip install openenv-core
openenv validate

License

MIT — built for Meta PyTorch OpenEnv Hackathon 2026.