Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 26 additions & 1 deletion CONTRIBUTING.md
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Contributing

Issues and pull requests are welcome. This guide covers local setup, the test suite, and how to add a task.
Issues and pull requests are welcome. This guide covers local setup, the test suite, adding a task, and deployment.

## Setup

Expand Down Expand Up @@ -50,6 +50,31 @@ The test suite checks:

When the reference query sorts on a key that can tie, keep in mind that any tie order counts as correct. Row order is not part of the correctness check.

## Deploying to Hugging Face Spaces

The live Space runs the root `Dockerfile`. Deploy with the OpenEnv CLI:

```bash
uv run openenv push
```

Hugging Face reads a Space's settings from a YAML header at the top of its `README.md`. The header is kept out of this repository's README because GitHub renders it as a table. `openenv push` adds a generic header automatically. To keep the project's name and description on the Space, put this block at the very top of the Space's `README.md`:

```yaml
---
title: SQL Query
emoji: 📊
colorFrom: blue
colorTo: indigo
sdk: docker
app_port: 8000
pinned: false
short_description: Text-to-SQL RL environment with execution-based rewards
tags:
- openenv
---
```

## Things to know

- **Import fallbacks.** The server runs in several layouts (`uv run server`, `uvicorn server.app:app`, `python -m sql_query_env.server.app`), and each one resolves imports differently. That's why several modules have `try`/`except ImportError` import blocks. Keep them.
Expand Down
226 changes: 128 additions & 98 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,45 +1,114 @@
---
title: SQL Query
emoji: 📊
colorFrom: blue
colorTo: indigo
sdk: docker
app_port: 8000
pinned: false
short_description: Text-to-SQL RL environment with execution-based rewards
tags:
- openenv
---
<h1 align="center">SQL Query</h1>

<p align="center">
A text-to-SQL reinforcement-learning environment for <a href="https://github.com/meta-pytorch/OpenEnv">OpenEnv</a>,<br>
with execution-based rewards and step-by-step diagnostic feedback.
</p>

<p align="center">
<a href="https://github.com/rajdeepchatale/sql_query_env/actions/workflows/ci.yml"><img src="https://github.com/rajdeepchatale/sql_query_env/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
<a href="LICENSE"><img src="https://img.shields.io/badge/license-BSD--3--Clause-blue.svg" alt="License: BSD-3-Clause"></a>
<img src="https://img.shields.io/badge/python-3.10%2B-blue.svg" alt="Python 3.10+">
<a href="https://huggingface.co/spaces/rajdeepchatale/sql_query_env"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Space-yellow.svg" alt="Hugging Face Space"></a>
</p>

<p align="center">
<a href="#quickstart">Quickstart</a> ·
<a href="#example-episode">Example</a> ·
<a href="#tasks">Tasks</a> ·
<a href="#reward">Reward</a> ·
<a href="https://huggingface.co/spaces/rajdeepchatale/sql_query_env">Live Space</a>
</p>

# SQL Query
---

[![CI](https://github.com/rajdeepchatale/sql_query_env/actions/workflows/ci.yml/badge.svg)](https://github.com/rajdeepchatale/sql_query_env/actions/workflows/ci.yml)
[![License: BSD-3-Clause](https://img.shields.io/badge/license-BSD--3--Clause-blue.svg)](LICENSE)
![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)
[![Hugging Face Space](https://img.shields.io/badge/%F0%9F%A4%97-Space-yellow.svg)](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 |
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -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

Expand Down
Loading