+ 
+ 
+ 
+
+ A text-to-SQL reinforcement-learning environment for OpenEnv,
+ with execution-based rewards and step-by-step diagnostic feedback.
+
+
+
+
+
+
+ Quickstart ยท + Example ยท + Tasks ยท + Reward ยท + Live Space +
-# SQL Query +--- -[](https://github.com/rajdeepchatale/sql_query_env/actions/workflows/ci.yml) -[](LICENSE) - -[](https://huggingface.co/spaces/rajdeepchatale/sql_query_env) +## Overview -**SQL Query** is an [OpenEnv](https://github.com/meta-pytorch/OpenEnv) reinforcement-learning environment for **text-to-SQL**. The agent receives a database schema and a question in plain English, then submits SQLite queries. Each query is executed and graded against a reference result. The reward gives partial credit for every part the query gets right, and each step also returns structured diagnostics that say what to fix next. +**SQL Query** trains and evaluates agents that turn natural-language questions into SQL. In each episode the agent receives a database schema and a question, then submits SQLite queries. Every query is executed against a sandboxed database and compared with a reference result. The agent gets a partial-credit reward and structured feedback on what to fix before its next attempt. -Built for the Meta PyTorch OpenEnv Hackathon 2026. +The environment was built for the Meta PyTorch OpenEnv Hackathon 2026. -## Why this environment +## Highlights -- **Verifiable answers.** Every task has a reference query. Grading executes the agent's query and compares the rows, so it is deterministic and needs no LLM judge. -- **Dense reward.** Five scoring components mean a query with the right tables but a wrong aggregate still earns signal, which binary pass/fail can't give. -- **Feedback an agent can act on.** Typed diagnostics such as `MISSING_JOIN`, `NULL_HANDLING` and `EXTRA_ROWS`, plus hints that get more specific after each failed attempt, support correction within an episode (process supervision). -- **Schemas to read, not memorize.** Three domains with different schemas and 14 tasks, from single-table filters to self-joins, anti-joins, and division-by-zero edge cases. -- **Safe to point an agent at.** The database is read-only, and every query runs under a time limit and a row limit. +- **Verifiable, deterministic grading.** Queries are executed and their rows compared with a reference result, with no LLM judge involved. +- **Dense reward signal.** Five scoring components give credit for partial progress, such as correct tables with a wrong aggregate. +- **Actionable feedback.** Typed diagnostics (`MISSING_JOIN`, `NULL_HANDLING`, `EXTRA_ROWS`, โฆ) and progressive hints support correction within an episode. +- **Generalization across schemas.** 14 tasks over three domains, ranging from single-table filters to self-joins, anti-joins, and division-by-zero edge cases. +- **Safe execution.** Read-only databases, per-query time and row limits, and a fresh database for every episode. -## How an episode works +## How it works -1. `reset()` picks a task, or `reset(task_id=...)` selects one, and builds a fresh in-memory SQLite database for its domain. +1. `reset()` selects a task, either at random or with `reset(task_id=...)`, and builds an in-memory SQLite database for its domain. 2. The observation contains the schema, the question, the expected output columns and row count, and a first hint. 3. The agent submits `SqlQueryAction(query="SELECT ...")`. 4. The grader executes the query, compares the result with the reference, and returns a reward in `[0, 1]` with feedback and diagnostics. -5. The episode ends when the result matches exactly or the task's attempt budget (5โ10) runs out. +5. The episode ends when the result matches exactly or the attempt budget (5โ10 per task) runs out. + +## Example episode + +Task `company_medium_2`, output from the environment: + +```text +Question Considering only active employees, find the average salary per department, + but only show departments where that average exceeds $100,000. ... +Expected columns: name, avg_salary ยท rows: 3 + +Attempt 1 SELECT d.name, AVG(e.salary) AS avg_salary + FROM employees e JOIN departments d ON e.department_id = d.id + GROUP BY d.name HAVING AVG(e.salary) > 100000 + ORDER BY avg_salary DESC + +Reward 0.82 (syntax 0.10 + tables 0.15 + columns 0.20 + results 0.30 + efficiency 0.07) +Diagnostic MISSING_ROWS: Got 2 rows but expected 3. + +Attempt 2 ... same query with WHERE e.is_active = 1 before GROUP BY ... + +Reward 0.97 Correct: the result set matches the expected output. Episode done. +``` + +## Quickstart + +Requires Python 3.10+ and [uv](https://docs.astral.sh/uv/). + +```bash +git clone https://github.com/rajdeepchatale/sql_query_env.git +cd sql_query_env +uv sync +uv run server # serves on http://localhost:8000 +``` + +Connect from Python: + +```python +import asyncio +from sql_query_env import SqlQueryAction, SqlQueryEnv + +async def main(): + async with SqlQueryEnv(base_url="http://localhost:8000") as env: + result = await env.reset(task_id="company_easy_1") + print(result.observation.question) + + result = await env.step(SqlQueryAction( + query="SELECT name, salary FROM employees ORDER BY salary DESC" + )) + print(result.reward, result.done) + for diag in result.observation.diagnostics: + print(f"[{diag['type']}] {diag['message']}") + +asyncio.run(main()) +``` + +Or run it in Docker: + +```bash +docker build -t sql_query_env . +docker run -p 8000:8000 sql_query_env +``` ## Tasks +Three database domains, each with its own schema: + | Domain | Tables | |---|---| | Company Analytics | departments, employees, products, customers, orders, reviews | @@ -73,15 +142,15 @@ Built for the Meta PyTorch OpenEnv Hackathon 2026. | Results | 0.45 | How closely the returned rows match the reference rows | | Efficiency | 0.10 | SQL style: aliases, explicit columns, no needless `DISTINCT` or comma joins | -Penalties: **โ0.10** for destructive SQL, which is also rejected without being executed, and **โ0.05** for resubmitting an identical query. The total is clamped to `[0, 1]`. +**Penalties:** โ0.10 for destructive SQL, which is also rejected without being executed, and โ0.05 for resubmitting an identical query. The total is clamped to `[0, 1]`. -How results are compared: +**Result comparison** -- Rows are compared as a **multiset**, so rows duplicated by a join fan-out count against the query. -- Numbers are compared by value to two decimal places, so `INTEGER 145000` equals `REAL 145000.0`. Text is compared case-insensitively. -- Full result credit requires an **exact match**. Row order does not affect correctness, because several reference queries sort on keys with ties. For partially correct results, rows already in the reference position earn a small bonus. -- A query counts as **correct** when its rows match exactly and all expected columns are present. Correctness ends the episode, not a score threshold. A correct query scores 0.92โ1.00 depending on its style points. -- Style bonuses that depend on the task, such as `COALESCE`, are only awarded when the reference query uses them. Sprinkling them into every query earns nothing. +- Rows are compared as a multiset, so duplicates from a join fan-out count against the query. +- Numbers are compared by value to two decimal places. Text is compared case-insensitively. +- Full result credit requires an exact match. Row order does not affect correctness, because several reference queries sort on keys with ties. Partially correct results earn a small bonus for rows already in the reference position. +- A query is **correct** when its rows match exactly and all expected columns are present. Correctness ends the episode. A correct query scores 0.92โ1.00, depending on its style points. +- Task-dependent style bonuses, such as `COALESCE`, are only awarded when the reference query uses them. ## Diagnostics @@ -116,63 +185,24 @@ Each step returns a list of `{type, severity, message, suggestion}` objects: | `expected_columns`, `expected_row_count` | Shape of the correct answer | | `query_result`, `query_error` | Formatted rows (first 20) or the SQLite error | | `feedback` | Score breakdown, quality notes, issues, and hints as text | -| `diagnostics`, `efficiency_notes` | Structured versions of the above | -| `steps_remaining`, `current_score`, `history` | Episode progress (best score so far, previous queries and scores) | +| `diagnostics`, `efficiency_notes` | Structured versions of the feedback | +| `steps_remaining`, `current_score`, `history` | Episode progress: best score so far, previous queries and scores | ## Sandboxing -- **Read-only database.** After seeding, an SQLite authorizer allows only read operations. SQLite itself refuses writes, schema changes, `PRAGMA` and `ATTACH`. -- **Resource limits.** Each query gets 2 seconds and at most 10,000 rows, so an unbounded recursive CTE can't block a server worker. +- **Read-only database.** After seeding, an SQLite authorizer permits only read operations. SQLite itself refuses writes, schema changes, `PRAGMA` and `ATTACH`. +- **Resource limits.** Each query gets 2 seconds and at most 10,000 rows, so runaway queries can't block a server worker. - **Isolation.** Every episode gets its own in-memory database, and each step accepts exactly one statement. -## Quickstart - -Requires Python 3.10+ and [uv](https://docs.astral.sh/uv/). - -```bash -git clone https://github.com/rajdeepchatale/sql_query_env.git -cd sql_query_env -uv sync -uv run server # serves on http://localhost:8000 -``` - -From Python: - -```python -import asyncio -from sql_query_env import SqlQueryAction, SqlQueryEnv - -async def main(): - async with SqlQueryEnv(base_url="http://localhost:8000") as env: - result = await env.reset(task_id="company_easy_1") - print(result.observation.question) - - result = await env.step(SqlQueryAction( - query="SELECT name, salary FROM employees ORDER BY salary DESC" - )) - print(result.reward, result.done) - for diag in result.observation.diagnostics: - print(f"[{diag['type']}] {diag['message']}") - -asyncio.run(main()) -``` - -With Docker: - -```bash -docker build -t sql_query_env . -docker run -p 8000:8000 sql_query_env -``` - ## Baseline agent -`inference.py` runs an OpenAI-compatible chat model on all 14 tasks. The default is `Qwen/Qwen2.5-72B-Instruct` via the Hugging Face router. The script logs each episode in the `[START]` / `[STEP]` / `[END]` format expected by the OpenEnv evaluation pipeline, then prints per-task and per-domain averages. +`inference.py` runs an OpenAI-compatible chat model on all 14 tasks. The default is `Qwen/Qwen2.5-72B-Instruct` via the Hugging Face router. It logs each episode in the `[START]` / `[STEP]` / `[END]` format expected by the OpenEnv evaluation pipeline, then prints per-task and per-domain averages. ```bash -export HF_TOKEN="hf_..." # or API_KEY for another provider -export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct" # optional -export API_BASE_URL="https://router.huggingface.co/v1" # optional -uv run python inference.py # expects the server on localhost:8000 +export HF_TOKEN="hf_..." # or API_KEY for another provider +export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct" # optional +export API_BASE_URL="https://router.huggingface.co/v1" # optional +uv run python inference.py # expects the server on localhost:8000 ``` Set `ENV_BASE_URL` to use a remote server, or `IMAGE_NAME` to start the environment from a Docker image. @@ -186,35 +216,35 @@ uv run ruff check . uv run openenv validate . ``` -See [CONTRIBUTING.md](CONTRIBUTING.md) for how to add a task. +See [CONTRIBUTING.md](CONTRIBUTING.md) for adding tasks and deploying to Hugging Face Spaces. ## Project structure -``` +```text sql_query_env/ -โโโ models.py # Action / observation Pydantic models -โโโ client.py # Async WebSocket client (EnvClient) -โโโ inference.py # LLM baseline agent -โโโ openenv.yaml # OpenEnv manifest: tasks, action and observation spaces -โโโ Dockerfile # Container image (HF Spaces, openenv build) +โโโ models.py # Action / observation Pydantic models +โโโ client.py # Async WebSocket client (EnvClient) +โโโ inference.py # LLM baseline agent +โโโ openenv.yaml # OpenEnv manifest: tasks, action and observation spaces +โโโ Dockerfile # Container image (HF Spaces, openenv build) โโโ server/ -โ โโโ app.py # FastAPI app and `server` entry point +โ โโโ app.py # FastAPI app and `server` entry point โ โโโ sql_query_env_environment.py # Episode logic: task selection, hints, termination -โ โโโ tasks.py # Schemas, seed data, read-only DB factory, task definitions -โ โโโ graders.py # Execution sandbox, 5-component scoring, diagnostics -โโโ tests/ # pytest suite +โ โโโ tasks.py # Schemas, seed data, read-only DB factory, tasks +โ โโโ graders.py # Execution sandbox, scoring, diagnostics +โโโ tests/ # pytest suite ``` ## Limitations -- **Diagnostics reveal the reference query's structure.** They name its tables and flag a missing `JOIN`, `GROUP BY` or `HAVING`. That is by design for guided correction, but it means the scores measure *assisted* text-to-SQL, not blind generation. +- **Guided, not blind.** Diagnostics reveal the reference query's structure: they name its tables and flag a missing `JOIN`, `GROUP BY` or `HAVING`. Scores therefore measure assisted text-to-SQL. - **Row order is not graded.** Tie-aware order checking would need an explicit sort specification per task. -- **Table matching is regex-based.** Table references are extracted without a full SQL parser, so CTE names are not resolved. -- **Small, synthetic data, SQLite dialect only.** Seed tables have 3โ24 rows each. +- **Regex-based table matching.** Table references are extracted without a full SQL parser, so CTE names are not resolved. +- **Small, synthetic data.** Seed tables have 3โ24 rows each, and only the SQLite dialect is supported. ## License -BSD 3-Clause. See [LICENSE](LICENSE). +Released under the [BSD 3-Clause License](LICENSE). ## Acknowledgements