Tsah00/sql-env

1---2title: SQL Query Learning Environment3emoji: 🗄️4colorFrom: blue5colorTo: green6sdk: docker7app_port: 78608tags:9  - openenv10  - reinforcement-learning11  - sql12  - agent13  - text-to-sql14license: mit15---16 17# SQL Query Learning Environment18 19An OpenEnv-compatible reinforcement learning environment that simulates the day-to-day work of a **data analyst** fulfilling ad-hoc SQL requests from business stakeholders (marketing, finance, CRM, merchandising). Agents learn to translate natural-language business requirements into correct, efficient SQL queries — a high-value real-world task.20 21## Why This Environment Matters22 23**Text-to-SQL is one of the most commercially valuable analyst skills** — data teams spend significant time writing bespoke queries for stakeholder requests, and current LLMs still fail on complex queries involving CTEs, window functions, and multi-table joins. Existing benchmarks (Spider, BIRD) are static evaluation sets. There is no standard RL environment for *training* agents to improve iteratively on SQL tasks through trial-and-error feedback.24 25This environment fills that gap:26 27- **Real-world task**: agents play the role of a data analyst, fulfilling concrete business requests — not solving academic puzzles28- **Iterative learning**: reward signal at every step, not just at episode end29- **Partial credit with exploit protection**: Jaccard-based row matching gives meaningful gradient signal while penalising both missing rows *and* extra rows — returning the entire table does not score 1.030- **Difficulty curriculum**: 3 tiers (easy → medium → hard) enable curriculum learning strategies31- **Business-realistic framing**: all 9 tasks are modelled on real analyst/engineer requests with named stakeholder teams32- **Zero external dependencies**: pure SQLite, runs on any 2 vCPU / 8 GB machine33 34## Environment Overview35 36The environment hosts an in-memory SQLite e-commerce database with 4 tables and deterministic seed data:37 38| Table | Rows | Key Columns |39|---|---|---|40| `customers` | 20 | id, name, email, city, country, created_at |41| `products` | 20 | id, name, category, price, stock |42| `orders` | 30 | id, customer_id, order_date, total_amount, status |43| `order_items` | 50 | id, order_id, product_id, quantity, unit_price |44 45At each step the agent submits a SQL query. The grader executes it, compares results to the reference solution, and returns a reward in `[0.0, 1.0]` with detailed feedback.46 47## Action Space48 49| Field | Type | Description |50|---|---|---|51| `query` | `str` | The SQL query string to execute |52| `task_id` | `str \| None` | Target a specific task (defaults to current) |53| `difficulty` | `str` | `"easy"`, `"medium"`, or `"hard"` |54 55## Observation Space56 57| Field | Type | Description |58|---|---|---|59| `result` | `list[dict]` | Rows returned by the agent's query |60| `error` | `str` | SQL error message (empty if none) |61| `reward` | `float` | Step reward in `[0.0, 1.0]` |62| `done` | `bool` | Whether the episode has ended |63| `message` | `str` | Grader feedback (e.g. "3/5 rows matched") |64| `schema_info` | `str` | Full database schema and role context |65| `task_description` | `str` | Business-context task description |66| `expected_columns` | `list[str]` | Column names the answer must contain |67| `step_count` | `int` | Steps taken this episode |68| `score_breakdown` | `dict` | Sub-scores: correctness, keyword_bonus, efficiency |69 70## Tasks (9 total across 3 tiers)71 72All tasks simulate real stakeholder requests submitted to a data analyst:73 74### Easy — Basic filtering and aggregation75 76| ID | Stakeholder Request | SQL Concepts |77|---|---|---|78| `easy_1` | Marketing: US customer email list for a promotional campaign | `WHERE` country filter |79| `easy_2` | Finance: daily fulfilment count report | `COUNT` + `WHERE` status |80| `easy_3` | Merchandising: top-5 premium products for a catalogue | `ORDER BY` + `LIMIT` |81 82### Medium — Joins, grouping, and date logic83 84| ID | Stakeholder Request | SQL Concepts |85|---|---|---|86| `medium_1` | Loyalty: total spend per customer for VIP tier selection | `JOIN` + `GROUP BY` + `SUM` |87| `medium_2` | Inventory: dead-stock products that have never been ordered | `LEFT JOIN` + `NULL` check |88| `medium_3` | Finance: monthly average order value trend for 2023 | `STRFTIME` + `AVG` + `GROUP BY` |89 90### Hard — CTEs, window functions, correlated subqueries91 92| ID | Stakeholder Request | SQL Concepts |93|---|---|---|94| `hard_1` | CRM: customers above average lifetime value for premium tier | `CTE` + scalar subquery |95| `hard_2` | Category mgmt: best-selling SKU per product category (ties allowed) | `CTE` + `RANK()` window function |96| `hard_3` | Retention: customers active in all three years 2022, 2023, and 2024 | Correlated subquery + `COUNT DISTINCT` |97 98## Reward Function99 100```101reward = correctness × 0.7 + keyword_bonus × 0.1 + efficiency_bonus × 0.2102```103 104| Component | Range | Description |105|---|---|---|106| `correctness` | 0.0–1.0 | **Jaccard multiset similarity** — penalises both missing rows (low recall) and extra rows (low precision). An agent cannot score 1.0 by returning the entire table. |107| `keyword_bonus` | 0.0–0.1 | Query uses expected SQL constructs (JOIN, GROUP BY, etc.) |108| `efficiency_bonus` | 0.0–0.2 | Penalises `SELECT *`, `CROSS JOIN`, and deeply nested subqueries |109 110**Why this reward design is good for RL:**111- Non-binary: agents receive gradient signal even on partially correct queries112- Exploit-resistant: Jaccard scoring penalises both over-fetching and under-fetching113- Column aliases normalised: `total` instead of `total_spent` is not penalised if values match114- Efficiency signal discourages degenerate solutions115- Per-step rewards enable policy gradient methods without sparse returns116 117**Grader variety (not always the same score):**118 119| Input | easy_1 reward |120|---|---|121| Correct `WHERE country = 'USA'` | 1.00 |122| Return all 20 customers (exploit attempt) | 0.34 |123| Empty result set | 0.30 |124| SQL syntax error | 0.00 |125 126## Baseline Scores127 128Deterministic rule-based baseline (no LLM) — achieved in 1 step per task:129 130| Task | Score | Steps |131|---|---|---|132| easy_1 | 1.00 | 1 |133| easy_2 | 1.00 | 1 |134| easy_3 | 1.00 | 1 |135| medium_1 | 1.00 | 1 |136| medium_2 | 1.00 | 1 |137| medium_3 | 1.00 | 1 |138| hard_1 | 1.00 | 1 |139| hard_2 | 1.00 | 1 |140| hard_3 | 1.00 | 1 |141 142LLM agents (e.g. Qwen2.5-72B, Nemotron) are expected to score **0.7–1.0 on easy**, **0.5–0.9 on medium**, and **0.2–0.7 on hard** — leaving meaningful room for RL improvement.143 144## Setup145 146### Local Development147 148```bash149pip install -r requirements.txt150 151# Start the server152uvicorn server.app:app --host 0.0.0.0 --port 7860153 154# Run inference (separate terminal)155python inference.py156```157 158### Docker159 160```bash161docker build -t sql_env .162docker run -p 7860:7860 \163  -e API_BASE_URL="https://router.huggingface.co/v1" \164  -e MODEL_NAME="Qwen/Qwen2.5-72B-Instruct" \165  -e HF_TOKEN="your-hf-token" \166  sql_env167```168 169### Run Tests170 171```bash172python -m pytest test_env.py -v173# 72 tests: database, graders, partial credit, lifecycle, all 9 tasks, FastAPI endpoints174```175 176## API Endpoints177 178| Method | Path | Description |179|---|---|---|180| `GET` | `/health` | Health check — returns HTTP 200 |181| `POST` | `/reset` | Start new episode |182| `POST` | `/step` | Execute SQL action |183| `GET` | `/state` | Current episode state |184| `GET` | `/tasks` | List all tasks by difficulty |185| `GET` | `/schema` | Database schema |186| `GET` | `/` | Interactive web UI |187| `WS` | `/ws` | WebSocket endpoint |188 189## Environment Variables190 191| Variable | Default | Description |192|---|---|---|193| `API_BASE_URL` | `https://router.huggingface.co/v1` | LLM API endpoint |194| `MODEL_NAME` | `Qwen/Qwen2.5-72B-Instruct` | Model identifier |195| `HF_TOKEN` / `OPENAI_API_KEY` | *(required)* | Hugging Face / API key |196 197The inference script accepts `HF_TOKEN`, `OPENAI_API_KEY`, or `API_KEY` (checked in that order).198 199## Project Structure200 201```202.203├── inference.py           # Baseline inference script (mandatory, root)204├── models.py              # Pydantic Action / Observation / State models205├── client.py              # HTTP EnvClient206├── openenv.yaml           # OpenEnv manifest207├── pyproject.toml         # Package metadata and dependencies208├── uv.lock                # uv lockfile for multi-mode deployment209├── requirements.txt       # pip dependencies210├── Dockerfile             # HF Spaces-compatible container211├── test_env.py            # 72-test suite212└── server/213    ├── app.py             # FastAPI server (reset / step / state / ws)214    ├── sql_environment.py # Core environment logic215    └── tasks.py           # 9 tasks, graders, SQLite seed data216```217 218## Judging Criteria Compliance219 220| Criterion | Status |221|---|---|222| Simulates real human task (data analyst SQL workflow) | ✓ |223| OpenEnv typed models (Action, Observation, State) | ✓ |224| `step()` / `reset()` / `state()` endpoints | ✓ |225| `openenv.yaml` with metadata | ✓ |226| 3+ tasks with programmatic graders | ✓ (9 tasks) |227| Easy → medium → hard difficulty range | ✓ |228| Deterministic graders with clear success/failure | ✓ |229| Partial-credit reward over full trajectory | ✓ |230| Exploit-resistant grader (Jaccard, no table-dump cheat) | ✓ |231| Baseline inference script with OpenAI client | ✓ |232| `HF_TOKEN` / `OPENAI_API_KEY` credential support | ✓ |233| `[START]` / `[STEP]` / `[END]` stdout log format | ✓ |234| Dockerfile builds (non-root, port 7860, HEALTHCHECK) | ✓ |235| `pyproject.toml` + `uv.lock` for multi-mode deployment | ✓ |236| Runs on 2 vCPU / 8 GB, inference < 20 min | ✓ |237