Skip to content

Commit 34c1564

Browse files
OfficialAbhinavSinghOfficialAbhinavSingh
committed
feat(sql-optim-env): add execution-grounded DuckDB SQL query optimization environment
1 parent 60ecc7b commit 34c1564

18 files changed

Lines changed: 2209 additions & 0 deletions

‎docs/source/_toctree.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,8 @@
131131
title: OpenCode
132132
- local: environments/sophistry_bench_sprint
133133
title: Sophistry Bench Sprint
134+
- local: environments/sql_optim
135+
title: SQL Query Optimization
134136
title: Environments
135137
- isExpanded: false
136138
sections:

‎docs/source/environments.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,13 @@ The OpenEnv community has built a catalog of ready-to-run environments that cove
265265
<a href="environments/sophistry_bench_sprint" class="!no-underline border dark:border-gray-700 px-3 py-1 rounded text-sm hover:shadow">📄 Docs</a>
266266
</div>
267267
</div>
268+
<div class="border dark:border-gray-700 p-5 rounded-lg shadow">
269+
<div class="font-bold mb-2">SQL Query Optimization</div>
270+
<p class="text-sm"><code>sql_optim_env</code> rewards an agent for rewriting slow SQL: it executes both the original and optimized query against an in-memory DuckDB seeded with synthetic data and scores measured speedup plus result-correctness across five anti-pattern tasks.</p>
271+
<div class="flex gap-2 mt-3">
272+
<a href="environments/sql_optim" class="!no-underline border dark:border-gray-700 px-3 py-1 rounded text-sm hover:shadow">📄 Docs</a>
273+
</div>
274+
</div>
268275
</div>
269276
</div>
270277

‎docs/source/environments/sql_optim.md‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
<!-- openenv-source: sql_optim_env -->
2+
# SQL Query Optimization Environment
3+
4+
An **execution-grounded** environment for training agents to write fast SQL. The
5+
agent is given a slow query plus its schema and must return a rewrite. Reward is
6+
**not** a keyword heuristic: the environment runs both the original and the
7+
optimized query against a real in-memory **DuckDB** database seeded with
8+
realistic synthetic data (10k users, 500k orders, 1k products, 1M events), then
9+
scores the agent on measured speedup and result-correctness.
10+
11+
Five tasks scale from basic anti-patterns to expert window-function audits.
12+
13+
## Quick Start
14+
15+
```python
16+
from sql_optim_env import SQLOptimAction, SQLOptimEnv
17+
18+
with SQLOptimEnv(base_url="http://localhost:8000").sync() as env:
19+
obs = env.reset(task_id="task_1_basic_antipatterns").observation
20+
print(obs.sql_query) # the slow query to fix
21+
print(obs.schema_info) # tables, row counts
22+
23+
result = env.step(
24+
SQLOptimAction(
25+
suggestions=[
26+
{"issue_type": "select_star", "line": 1,
27+
"description": "SELECT * reads every column", "severity": "high",
28+
"fix": "project only the needed columns"},
29+
],
30+
optimized_query="SELECT id FROM orders WHERE customer_id = 5000",
31+
summary="Dropped SELECT * and the non-sargable CAST.",
32+
estimated_improvement="~10x faster",
33+
approved=False,
34+
)
35+
)
36+
print(result.reward, result.done)
37+
print(result.metadata["reward_breakdown"])
38+
print(result.metadata["execution"]["speedup"]) # real DuckDB speedup
39+
```
40+
41+
Async usage mirrors the other environments:
42+
43+
```python
44+
async with SQLOptimEnv(base_url="http://localhost:8000") as env:
45+
result = await env.reset(task_id="task_1_basic_antipatterns")
46+
```
47+
48+
## Action
49+
50+
| Field | Type | Meaning |
51+
|-------|------|---------|
52+
| `suggestions` | `list[dict]` | Issues found: `issue_type`, `line`, `description`, `severity`, `fix` |
53+
| `optimized_query` | `str` | The rewritten SQL — executed for real |
54+
| `summary` | `str` | Overall analysis |
55+
| `estimated_improvement` | `str` | Agent's own speedup estimate |
56+
| `approved` | `bool` | `True` if the query is already optimal |
57+
58+
## Observation
59+
60+
`sql_query`, `schema_info`, `task_description`, `difficulty`, `step_count`,
61+
`max_steps`, `issues_found_so_far`, and `last_execution` (the previous step's
62+
real timing comparison, so the agent can refine its rewrite). The step reward is
63+
on `reward`; the composite `reward_breakdown`, `feedback`, and full `execution`
64+
comparison are under `metadata`.
65+
66+
## Reward
67+
68+
Composite score in `[0, 1]`, computed inside the environment:
69+
70+
| Component | Weight | Source |
71+
|-----------|--------|--------|
72+
| Execution speedup | 35% | real DuckDB timing ratio |
73+
| Result correctness | 20% | do both queries return identical rows? |
74+
| Issue detection | 25% | flagged issues vs. ground truth |
75+
| Approval correctness | 8% | correctly judged the query bad/good |
76+
| Summary quality | 7% | thoroughness of the written analysis |
77+
| Severity labels | 5% | severities present |
78+
79+
## Tasks
80+
81+
Five tasks (`task_1_basic_antipatterns` … `task_5_*`) covering `SELECT *`,
82+
non-sargable predicates, functions on filter columns, join ordering, and window
83+
functions, spanning `easy` → `expert`.
84+
85+
## Run the server locally
86+
87+
```bash
88+
uvicorn server.app:app --host 0.0.0.0 --port 8000
89+
```
90+
91+
## Tests
92+
93+
```bash
94+
PYTHONPATH=src:envs uv run pytest tests/envs/test_sql_optim_environment.py -v
95+
```

‎envs/sql_optim_env/README.md‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
---
2+
title: SQL Query Optimization Environment
3+
emoji: 🗄️
4+
colorFrom: indigo
5+
colorTo: blue
6+
sdk: docker
7+
pinned: false
8+
app_port: 8000
9+
base_path: /web
10+
tags:
11+
- openenv
12+
- sql
13+
- duckdb
14+
- optimization
15+
---
16+
17+
# SQL Query Optimization Environment
18+
19+
An **execution-grounded** environment for training agents to write fast SQL. The
20+
agent is given a slow query plus its schema and must return a rewrite. Reward is
21+
**not** a keyword heuristic: the environment runs both the original and the
22+
optimized query against a real in-memory **DuckDB** database seeded with
23+
realistic synthetic data (10k users, 500k orders, 1k products, 1M events), then
24+
scores the agent on measured speedup and result-correctness.
25+
26+
Five tasks scale from basic anti-patterns to expert window-function audits.
27+
28+
## Quick Start
29+
30+
```python
31+
from sql_optim_env import SQLOptimAction, SQLOptimEnv
32+
33+
with SQLOptimEnv(base_url="http://localhost:8000").sync() as env:
34+
obs = env.reset(task_id="task_1_basic_antipatterns").observation
35+
print(obs.sql_query) # the slow query to fix
36+
print(obs.schema_info) # tables, row counts
37+
38+
result = env.step(
39+
SQLOptimAction(
40+
suggestions=[
41+
{"issue_type": "select_star", "line": 1,
42+
"description": "SELECT * reads every column", "severity": "high",
43+
"fix": "project only the needed columns"},
44+
],
45+
optimized_query="SELECT id FROM orders WHERE customer_id = 5000",
46+
summary="Dropped SELECT * and the non-sargable CAST.",
47+
estimated_improvement="~10x faster",
48+
approved=False,
49+
)
50+
)
51+
print(result.reward, result.done)
52+
print(result.metadata["reward_breakdown"])
53+
print(result.metadata["execution"]["speedup"]) # real DuckDB speedup
54+
```
55+
56+
Async usage mirrors the other environments:
57+
58+
```python
59+
async with SQLOptimEnv(base_url="http://localhost:8000") as env:
60+
result = await env.reset(task_id="task_1_basic_antipatterns")
61+
```
62+
63+
## Action
64+
65+
| Field | Type | Meaning |
66+
|-------|------|---------|
67+
| `suggestions` | `list[dict]` | Issues found: `issue_type`, `line`, `description`, `severity`, `fix` |
68+
| `optimized_query` | `str` | The rewritten SQL — executed for real |
69+
| `summary` | `str` | Overall analysis |
70+
| `estimated_improvement` | `str` | Agent's own speedup estimate |
71+
| `approved` | `bool` | `True` if the query is already optimal |
72+
73+
## Observation
74+
75+
`sql_query`, `schema_info`, `task_description`, `difficulty`, `step_count`,
76+
`max_steps`, `issues_found_so_far`, and `last_execution` (the previous step's
77+
real timing comparison, so the agent can refine its rewrite). The step reward is
78+
on `reward`; the composite `reward_breakdown`, `feedback`, and full `execution`
79+
comparison are under `metadata`.
80+
81+
## Reward
82+
83+
Composite score in `[0, 1]`, computed inside the environment:
84+
85+
| Component | Weight | Source |
86+
|-----------|--------|--------|
87+
| Execution speedup | 35% | real DuckDB timing ratio |
88+
| Result correctness | 20% | do both queries return identical rows? |
89+
| Issue detection | 25% | flagged issues vs. ground truth |
90+
| Approval correctness | 8% | correctly judged the query bad/good |
91+
| Summary quality | 7% | thoroughness of the written analysis |
92+
| Severity labels | 5% | severities present |
93+
94+
## Tasks
95+
96+
Five tasks (`task_1_basic_antipatterns` … `task_5_*`) covering `SELECT *`,
97+
non-sargable predicates, functions on filter columns, join ordering, and window
98+
functions, spanning `easy` → `expert`.
99+
100+
## Run the server locally
101+
102+
```bash
103+
uvicorn server.app:app --host 0.0.0.0 --port 8000
104+
```
105+
106+
## Tests
107+
108+
```bash
109+
PYTHONPATH=src:envs uv run pytest tests/envs/test_sql_optim_environment.py -v
110+
```

‎envs/sql_optim_env/__init__.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
"""
8+
SQL Query Optimization Environment for OpenEnv.
9+
10+
An execution-grounded RL environment: the agent rewrites slow SQL and is
11+
rewarded by real DuckDB execution speedup plus result-correctness across five
12+
tasks that scale from basic anti-patterns to expert window-function audits.
13+
14+
Examples:
15+
16+
```python
17+
from envs.sql_optim_env import SQLOptimEnv, SQLOptimAction
18+
19+
with SQLOptimEnv(base_url="http://localhost:8000").sync() as env:
20+
obs = env.reset(task_id="task_1_basic_antipatterns").observation
21+
result = env.step(SQLOptimAction(optimized_query="SELECT id FROM orders ..."))
22+
print(result.reward, result.done)
23+
```
24+
"""
25+
26+
from .client import SQLOptimEnv
27+
from .models import SQLOptimAction, SQLOptimObservation, SQLOptimState
28+
29+
__all__ = [
30+
"SQLOptimEnv",
31+
"SQLOptimAction",
32+
"SQLOptimObservation",
33+
"SQLOptimState",
34+
]

‎envs/sql_optim_env/client.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
"""
8+
Client for the SQL Query Optimization environment.
9+
10+
Maintains a persistent WebSocket connection to the environment server for
11+
efficient multi-step episodes. Async by default; use `.sync()` for a
12+
synchronous wrapper.
13+
"""
14+
15+
from __future__ import annotations
16+
17+
from typing import Any, Dict
18+
19+
from openenv.core.client_types import StepResult
20+
from openenv.core.env_client import EnvClient
21+
22+
from .models import SQLOptimAction, SQLOptimObservation, SQLOptimState
23+
24+
25+
class SQLOptimEnv(EnvClient[SQLOptimAction, SQLOptimObservation, SQLOptimState]):
26+
"""
27+
Client for the SQL Query Optimization environment.
28+
29+
The agent receives a slow SQL query plus its schema, returns a
30+
[`SQLOptimAction`][envs.sql_optim_env.models.SQLOptimAction] with a
31+
rewritten `optimized_query`, and is rewarded by real DuckDB execution
32+
speedup and result-correctness.
33+
34+
Examples:
35+
36+
```python
37+
with SQLOptimEnv(base_url="http://localhost:8000").sync() as env:
38+
obs = env.reset(task_id="task_1_basic_antipatterns").observation
39+
result = env.step(
40+
SQLOptimAction(
41+
optimized_query="SELECT id FROM orders WHERE customer_id = 5000",
42+
summary="Dropped SELECT * and the non-sargable CAST.",
43+
)
44+
)
45+
print(result.reward, result.done)
46+
```
47+
"""
48+
49+
def _step_payload(self, action: SQLOptimAction) -> Dict[str, Any]:
50+
"""Serialize a [`SQLOptimAction`] into the step request payload."""
51+
return {
52+
"suggestions": action.suggestions,
53+
"optimized_query": action.optimized_query,
54+
"summary": action.summary,
55+
"estimated_improvement": action.estimated_improvement,
56+
"approved": action.approved,
57+
}
58+
59+
def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SQLOptimObservation]:
60+
"""Parse a server step/reset response into a `StepResult`."""
61+
obs_data = payload.get("observation", {})
62+
observation = SQLOptimObservation(**obs_data)
63+
return StepResult(
64+
observation=observation,
65+
reward=payload.get("reward", observation.reward),
66+
done=payload.get("done", observation.done),
67+
metadata=payload.get("metadata", observation.metadata),
68+
)
69+
70+
def _parse_state(self, payload: Dict[str, Any]) -> SQLOptimState:
71+
"""Parse a server state response into a [`SQLOptimState`]."""
72+
return SQLOptimState(**payload)

0 commit comments

Comments
 (0)