SQLArenaEnv - multi-step SQL reasoning environment

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,11 +1,268 @@
 ---
-title: Sql Arena Env
-emoji: 🐨
-colorFrom: indigo
-colorTo: red
+title: SQLArenaEnv
 sdk: docker
-pinned: false
-license: mit
+pinned: true
+tags:
+  - openenv
+  - rl-environment
+  - sql
+  - reasoning
+  - multi-step
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SQLArenaEnv
+
+> **The first OpenEnv-compatible environment for training multi-step SQL reasoning agents — where exploration queries are first-class actions.**
+
+[![OpenEnv](https://img.shields.io/badge/OpenEnv-compatible-blue)](https://github.com/meta-pytorch/OpenEnv)
+[![Tasks](https://img.shields.io/badge/tasks-50-green)]()
+[![Difficulty](https://img.shields.io/badge/difficulty-easy→expert-orange)]()
+
+---
+
+## What Is SQLArenaEnv?
+
+Most SQL benchmarks test single-shot generation — the model sees a question and must output the correct query in one attempt. Real-world SQL reasoning doesn't work that way. Analysts explore the data, run investigative queries, check schemas, and refine their approach before committing to a final answer.
+
+**SQLArenaEnv makes exploration a first-class action.**
+
+The agent is given a natural language question and a database. It can run up to 5 **explore** queries to investigate the schema and data — seeing real results — before submitting a final **submit** query that is scored. Agents that explore strategically outperform agents that blindly guess.
+
+This tests the skill that actually matters: **reasoning about data, not just memorizing SQL syntax.**
+
+---
+
+## The Core Mechanic
+
+```
+Episode start
+     │
+     ▼
+┌─────────────────────────────────────┐
+│  Question: "Find customers who      │
+│  placed more than 1 completed order"│
+│  Schema: customers, orders,         │
+│          order_items                │
+└──────────────┬──────────────────────┘
+               │
+    ┌──────────▼──────────┐
+    │   EXPLORE action     │  ← up to 5 free exploration queries
+    │   sql = "SELECT *    │    each returns real data rows
+    │   FROM customers     │    small -0.02 cost per step
+    │   LIMIT 5"           │    to discourage random fishing
+    └──────────┬───────────┘
+               │  (repeat up to 5 times)
+               │
+    ┌──────────▼──────────┐
+    │   SUBMIT action      │  ← final answer query
+    │   sql = "SELECT      │    scored against reference solution
+    │   c.name, COUNT(*)   │    1.0 = correct
+    │   FROM customers c   │    0.4 = partial
+    │   JOIN orders o ...  │    0.0 = wrong
+    │   HAVING COUNT > 1"  │   -0.1 = syntax error
+    └─────────────────────┘
+```
+
+---
+
+## Task Library — 50 Tasks Across 4 Tiers
+
+| Tier | Count | SQL Concepts Tested |
+|------|-------|---------------------|
+| **Easy** | 10 | SELECT, WHERE, ORDER BY, GROUP BY, basic aggregation |
+| **Medium** | 15 | JOINs, HAVING, subqueries, LEFT JOIN, correlated filters |
+| **Hard** | 15 | CTEs, window functions (RANK, LAG, PERCENT_RANK), multi-join analytics |
+| **Expert** | 10 | Correlated subqueries, financial scoring, multi-CTE chains, complex aggregation |
+
+All tasks use realistic domains: **e-commerce orders, HR systems, retail analytics, banking/finance.** Data is India-relevant (names, cities, currencies in INR).
+
+---
+
+## Reward Structure
+
+```python
+# Explore step
+reward = -0.02   # small cost — discourages blind fishing
+
+# Submit step
+reward = 1.0     # correct answer (exact match)
+reward = 0.4     # partial credit (right columns, wrong rows)
+reward = 0.0     # wrong answer
+reward = -0.1    # SQL syntax error
+```
+
+---
+
+## Quick Start
+
+### Install client
+```bash
+pip install git+https://huggingface.co/spaces/sakthivarshans/sql-arena-env
+```
+
+### Use in Python (async)
+```python
+import asyncio
+from sql_arena_env import SQLArenaEnv, SQLArenaAction
+
+async def main():
+    async with SQLArenaEnv(base_url="https://sakthivarshans-sql-arena-env.hf.space") as env:
+
+        # Start episode — random task
+        result = await env.reset()
+        obs = result.observation
+        print(f"Question: {obs.question}")
+        print(f"Schema:   {obs.schema_info}")
+
+        # Explore the data
+        result = await env.step(SQLArenaAction(
+            sql="SELECT * FROM customers LIMIT 5",
+            query_type="explore"
+        ))
+        print(f"Sample data: {result.observation.query_result}")
+
+        # Submit final answer
+        result = await env.step(SQLArenaAction(
+            sql="SELECT c.name, COUNT(*) as order_count FROM customers c JOIN orders o ON c.customer_id = o.customer_id WHERE o.status='completed' GROUP BY c.customer_id HAVING COUNT(*) > 1",
+            query_type="submit"
+        ))
+        print(f"Correct: {result.observation.is_correct}")
+        print(f"Reward:  {result.reward}")
+        print(f"Feedback: {result.observation.feedback}")
+
+asyncio.run(main())
+```
+
+### Load a specific task
+```python
+result = await env.reset(task_id="hard_002")   # specific task
+result = await env.reset(difficulty="medium")   # random from tier
+result = await env.reset()                       # fully random
+```
+
+### Sync usage
+```python
+with SQLArenaEnv(base_url="http://localhost:8000").sync() as env:
+    result = env.reset(task_id="easy_001")
+    result = env.step(SQLArenaAction(sql="SELECT 1", query_type="submit"))
+```
+
+---
+
+## HTTP API Reference
+
+| Method | Endpoint | Description |
+|--------|----------|-------------|
+| `GET` | `/health` | Liveness check |
+| `POST` | `/reset` | Start new episode |
+| `POST` | `/step` | Execute SQL action |
+| `GET` | `/state` | Current episode state |
+| `GET` | `/schema` | Action/observation schema |
+| `WS` | `/ws` | WebSocket (persistent session, use for training) |
+
+### Reset request body
+```json
+{ "task_id": "medium_001" }
+{ "difficulty": "hard" }
+{}
+```
+
+### Step request body
+```json
+{
+  "action": {
+    "sql": "SELECT name FROM customers LIMIT 5",
+    "query_type": "explore"
+  }
+}
+```
+
+---
+
+## Graders
+
+Four deterministic graders, one per difficulty tier. Each runs 3 representative tasks.
+
+```bash
+python graders.py
+```
+
+```
+easy      : 1.0000  ████████████████████
+medium    : 1.0000  ████████████████████
+hard      : 1.0000  ████████████████████
+expert    : 1.0000  ████████████████████
+overall   : 1.0000  ████████████████████
+```
+
+---
+
+## Run Locally
+
+### With Docker
+```bash
+docker build -t sql-arena-env:latest -f server/Dockerfile .
+docker run -p 8000:8000 sql-arena-env:latest
+curl http://localhost:8000/health
+```
+
+### Without Docker
+```bash
+pip install openenv-core
+pip install -e .
+uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+### Run graders
+```bash
+python graders.py
+```
+
+### Run inference script
+```bash
+export API_BASE_URL="https://router.huggingface.co/v1"
+export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct"
+export HF_TOKEN="hf_..."
+export SQLARENA_TASK="medium_001"
+python inference.py
+```
+
+---
+
+## Project Structure
+
+```
+sql_arena_env/
+├── __init__.py              # Package exports
+├── models.py                # SQLArenaAction, SQLArenaObservation
+├── client.py                # Async typed client (SQLArenaEnv)
+├── tasks.py                 # 50 curated SQL tasks with schemas & solutions
+├── graders.py               # 4 deterministic graders
+├── inference.py             # Hackathon inference script
+├── openenv.yaml             # OpenEnv manifest
+├── pyproject.toml           # pip installable
+└── server/
+    ├── app.py               # FastAPI + WebSocket via create_app()
+    ├── sql_arena_environment.py  # Core environment logic
+    ├── Dockerfile           # openenv-base multi-stage build
+    └── requirements.txt
+```
+
+---
+
+## Why SQLArenaEnv?
+
+**The gap it fills:** Text-to-SQL benchmarks like Spider and BIRD measure single-shot accuracy. No existing OpenEnv environment measures *multi-step SQL reasoning* where the agent can gather information before committing. This is the benchmark that matches how SQL is actually used.
+
+**Why exploration matters for RL training:** An agent that learns to run `SELECT * FROM table LIMIT 5` before attempting a complex GROUP BY query is learning a genuinely useful cognitive strategy — the same strategy a senior data analyst uses. Standard single-shot SQL environments cannot teach this. SQLArenaEnv can.
+
+**What improves with training:** GRPO/PPO agents trained on SQLArenaEnv learn to use explore steps strategically — they converge to running schema-discovery queries first (`SELECT * FROM sqlite_master`), then sample queries, then submitting. This mirrors expert human behavior and transfers to real SQL tasks.
+
+---
+
+## Citation
+
+```
+SQLArenaEnv — Multi-step SQL Reasoning Environment for OpenEnv
+OpenEnv Hackathon 2026 — Meta × Hugging Face × Scaler
+```

__init__.py

@@ -0,0 +1,6 @@
+"""SQLArenaEnv — Multi-step SQL reasoning environment for OpenEnv."""
+
+from .models import SQLArenaAction, SQLArenaObservation
+from .client import SQLArenaEnv
+
+__all__ = ["SQLArenaAction", "SQLArenaObservation", "SQLArenaEnv"]

client.py

@@ -0,0 +1,101 @@
+"""
+SQLArenaEnv Client.
+
+Typed async client for SQLArenaEnv. Maintains a persistent WebSocket
+connection for efficient multi-step interactions during RL training.
+
+Example (async):
+    async with SQLArenaEnv(base_url="https://your-space.hf.space") as env:
+        result = await env.reset(task_id="medium_001")
+        print(result.observation.question)
+
+        result = await env.step(SQLArenaAction(
+            sql="SELECT * FROM customers LIMIT 5",
+            query_type="explore"
+        ))
+        print(result.observation.query_result)
+
+        result = await env.step(SQLArenaAction(
+            sql="SELECT name, city FROM customers WHERE segment='Corporate'",
+            query_type="submit"
+        ))
+        print(result.observation.feedback)
+
+Example (sync):
+    with SQLArenaEnv(base_url="http://localhost:8000").sync() as env:
+        result = env.reset(difficulty="easy")
+        result = env.step(SQLArenaAction(sql="SELECT 1", query_type="submit"))
+
+Example (from Docker):
+    env = await SQLArenaEnv.from_docker_image("sql-arena-env:latest")
+    try:
+        result = await env.reset()
+    finally:
+        await env.close()
+"""
+
+from typing import Dict, Optional
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from .models import SQLArenaAction, SQLArenaObservation
+
+
+class SQLArenaEnv(EnvClient[SQLArenaAction, SQLArenaObservation, State]):
+    """
+    Async client for SQLArenaEnv.
+
+    This client maintains a persistent WebSocket connection to the environment
+    server, enabling efficient multi-step SQL interactions with low latency.
+    """
+
+    def _step_payload(self, action: SQLArenaAction) -> Dict:
+        return {
+            "sql": action.sql,
+            "query_type": action.query_type,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[SQLArenaObservation]:
+        obs_data = payload.get("observation", {})
+        observation = SQLArenaObservation(
+            # Task context
+            task_id=obs_data.get("task_id", ""),
+            difficulty=obs_data.get("difficulty", "easy"),
+            question=obs_data.get("question", ""),
+            schema_info=obs_data.get("schema_info", ""),
+
+            # Query result
+            query_result=obs_data.get("query_result", []),
+            query_error=obs_data.get("query_error"),
+            query_type=obs_data.get("query_type", "explore"),
+            rows_returned=obs_data.get("rows_returned", 0),
+
+            # Episode progress
+            explore_steps_used=obs_data.get("explore_steps_used", 0),
+            explore_steps_remaining=obs_data.get("explore_steps_remaining", 5),
+            submitted=obs_data.get("submitted", False),
+
+            # Feedback after submit
+            is_correct=obs_data.get("is_correct"),
+            feedback=obs_data.get("feedback"),
+            expected_row_count=obs_data.get("expected_row_count"),
+
+            # Base fields
+            done=payload.get("done", False),
+            reward=payload.get("reward", 0.0),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward", 0.0),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )
+        

graders.py

@@ -0,0 +1,86 @@
+"""
+SQLArenaEnv Graders.
+
+Four deterministic graders, one per difficulty tier.
+Each grader runs 3 fixed tasks using the reference solution agent
+and returns a score in [0.0, 1.0].
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from server.sql_arena_environment import SQLArenaEnvironment
+from models import SQLArenaAction
+from tasks import get_task
+
+
+def _grade_task(task_id: str) -> float:
+    env = SQLArenaEnvironment()
+    task = get_task(task_id)
+    if task is None:
+        return 0.0
+    env.reset(task_id=task_id)
+    obs = env.step(SQLArenaAction(sql=task.solution_sql, query_type="submit"))
+    return float(obs.reward) if obs.reward is not None else 0.0
+
+
+def _grade_tier(task_ids: list, tier_name: str) -> float:
+    scores = []
+    for tid in task_ids:
+        score = _grade_task(tid)
+        scores.append(score)
+        print(f"  {tid}: {score:.4f}")
+    mean = round(sum(scores) / len(scores), 4)
+    print(f"  {tier_name} mean: {mean:.4f}")
+    return mean
+
+
+def grade_easy() -> float:
+    """Easy grader — single table SELECT, WHERE, GROUP BY."""
+    print("Running easy grader...")
+    return _grade_tier(["easy_001", "easy_002", "easy_005"], "easy")
+
+
+def grade_medium() -> float:
+    """Medium grader — JOINs, HAVING, subqueries."""
+    print("Running medium grader...")
+    return _grade_tier(["medium_001", "medium_004", "medium_007"], "medium")
+
+
+def grade_hard() -> float:
+    """Hard grader — CTEs, window functions, complex aggregation."""
+    print("Running hard grader...")
+    return _grade_tier(["hard_001", "hard_003", "hard_006"], "hard")
+
+
+def grade_expert() -> float:
+    """Expert grader — correlated subqueries, financial analytics."""
+    print("Running expert grader...")
+    return _grade_tier(["expert_001", "expert_003", "expert_007"], "expert")
+
+
+def run_all_graders() -> dict:
+    """Run all four graders and return summary dict."""
+    results = {
+        "easy":   grade_easy(),
+        "medium": grade_medium(),
+        "hard":   grade_hard(),
+        "expert": grade_expert(),
+    }
+    results["overall"] = round(sum(v for k, v in results.items() if k != "overall") / 4, 4)
+    return results
+
+
+if __name__ == "__main__":
+    print("=" * 50)
+    print("SQLArenaEnv Graders")
+    print("=" * 50)
+    results = run_all_graders()
+    print()
+    print("=" * 50)
+    print("SUMMARY")
+    print("=" * 50)
+    for k, v in results.items():
+        bar = "█" * int(v * 20)
+        print(f"  {k:10s}: {v:.4f}  {bar}")

inference.py

@@ -0,0 +1,213 @@
+"""
+SQLArenaEnv Inference Script.
+
+Runs an LLM agent against SQLArenaEnv, emitting the mandatory
+[START] / [STEP] / [END] log format required by the hackathon evaluator.
+
+The agent strategy:
+  1. Read the question and schema from the reset observation
+  2. Run 1-2 exploratory queries to understand the data
+  3. Submit a final SQL query as the answer
+
+Usage:
+    API_BASE_URL=<url> MODEL_NAME=<model> HF_TOKEN=<token> python inference.py
+"""
+
+import asyncio
+import os
+import textwrap
+from typing import List, Optional
+
+from openai import OpenAI
+from client import SQLArenaEnv, SQLArenaAction
+
+IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") or os.getenv("IMAGE_NAME")
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+TASK_NAME = os.getenv("SQLARENA_TASK", "medium_001")
+BENCHMARK = "sql_arena_env"
+MAX_STEPS = 8
+TEMPERATURE = 0.2
+MAX_TOKENS = 512
+SUCCESS_SCORE_THRESHOLD = 0.5
+
+SYSTEM_PROMPT = textwrap.dedent("""
+You are an expert SQL agent working with a SQLite database.
+You will be given a natural language question and a schema description.
+Your job is to write correct SQLite SQL to answer the question.
+
+Rules:
+- Use only standard SQLite syntax (no MySQL/PostgreSQL specific features)
+- You can run EXPLORE queries first to understand the data structure
+- When ready, submit your final answer with query_type="submit"
+- SQLite date functions: SUBSTR(date,1,7) for YYYY-MM, julianday() for date math
+- Window functions available: ROW_NUMBER(), RANK(), DENSE_RANK(), LAG(), LEAD(), SUM() OVER(), AVG() OVER(), PERCENT_RANK()
+
+Response format — respond with ONLY a JSON object like this:
+{"sql": "SELECT ...", "query_type": "explore"}
+or
+{"sql": "SELECT ...", "query_type": "submit"}
+
+No explanation, no markdown, just the JSON.
+""").strip()
+
+
+def log_start(task: str, env: str, model: str) -> None:
+    print(f"[START] task={task} env={env} model={model}", flush=True)
+
+
+def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
+    error_val = error if error else "null"
+    done_val = str(done).lower()
+    action_safe = action.replace("\n", " ").replace("\r", "")[:200]
+    print(
+        f"[STEP] step={step} action={action_safe} reward={reward:.2f} done={done_val} error={error_val}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
+    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+    print(
+        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
+        flush=True,
+    )
+
+
+def build_user_prompt(obs, step: int, history: List[str]) -> str:
+    history_block = "\n".join(history[-3:]) if history else "None"
+    result_str = str(obs.query_result[:5]) if obs.query_result else "No results yet"
+
+    return textwrap.dedent(f"""
+    QUESTION: {obs.question}
+    SCHEMA: {obs.schema_info}
+
+    Step: {step}
+    Explore steps remaining: {obs.explore_steps_remaining}
+    Last query result (first 5 rows): {result_str}
+    Last error: {obs.query_error or 'None'}
+
+    Previous actions:
+    {history_block}
+
+    {"⚠️ No more explore steps — you MUST submit now (query_type='submit')" if obs.explore_steps_remaining == 0 else "You can explore more or submit your final answer."}
+
+    Respond with ONLY a JSON object: {{"sql": "...", "query_type": "explore" or "submit"}}
+    """).strip()
+
+
+def get_model_action(client: OpenAI, obs, step: int, history: List[str]):
+    """Call the LLM and parse its response into a SQLArenaAction."""
+    import json
+
+    user_prompt = build_user_prompt(obs, step, history)
+    try:
+        completion = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ],
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+        )
+        text = (completion.choices[0].message.content or "").strip()
+
+        # Strip markdown if present
+        if text.startswith("```"):
+            text = text.split("```")[1]
+            if text.startswith("json"):
+                text = text[4:]
+        text = text.strip()
+
+        parsed = json.loads(text)
+        sql = parsed.get("sql", "SELECT 1")
+        query_type = parsed.get("query_type", "explore")
+
+        # Force submit if no explore steps left
+        if obs.explore_steps_remaining == 0:
+            query_type = "submit"
+
+        return SQLArenaAction(sql=sql, query_type=query_type)
+
+    except Exception as exc:
+        print(f"[DEBUG] Model parse error: {exc} | raw: {text[:200] if 'text' in dir() else 'N/A'}", flush=True)
+        # Fallback: try a simple submit
+        return SQLArenaAction(
+            sql=f"SELECT * FROM sqlite_master WHERE type='table'",
+            query_type="explore" if obs.explore_steps_remaining > 0 else "submit"
+        )
+
+
+async def main() -> None:
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+
+    # Connect to environment
+    if IMAGE_NAME:
+        env = await SQLArenaEnv.from_docker_image(IMAGE_NAME)
+    else:
+        hf_space_url = os.getenv("HF_SPACE_URL", "http://localhost:8000")
+        env = SQLArenaEnv(base_url=hf_space_url)
+
+    history: List[str] = []
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+
+    try:
+        result = await env.reset(task_id=TASK_NAME)
+        obs = result.observation
+
+        for step in range(1, MAX_STEPS + 1):
+            if result.done:
+                break
+
+            action = get_model_action(client, obs, step, history)
+            result = await env.step(action)
+            obs = result.observation
+
+            reward = result.reward or 0.0
+            done = result.done
+            error = obs.query_error
+
+            rewards.append(reward)
+            steps_taken = step
+
+            log_step(
+                step=step,
+                action=f"{action.query_type}:{action.sql[:100]}",
+                reward=reward,
+                done=done,
+                error=error,
+            )
+
+            history.append(
+                f"Step {step} [{action.query_type}]: {action.sql[:80]} → "
+                f"rows={obs.rows_returned} reward={reward:.2f}"
+            )
+
+            if done:
+                break
+
+        # Score: correctness of final submission (1.0 = correct, 0.4 = partial, 0.0 = wrong)
+        final_reward = rewards[-1] if rewards else 0.0
+        score = final_reward  # already normalized [0,1]
+        success = score >= SUCCESS_SCORE_THRESHOLD
+
+    except Exception as exc:
+        print(f"[DEBUG] Episode error: {exc}", flush=True)
+
+    finally:
+        try:
+            await env.close()
+        except Exception as e:
+            print(f"[DEBUG] env.close() error: {e}", flush=True)
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

models.py

@@ -0,0 +1,70 @@
+"""
+Data models for SQLArenaEnv.
+
+SQLArenaEnv is a multi-step SQL reasoning environment where agents must
+explore a database schema, run investigative queries, and submit a final
+answer — testing true SQL understanding, not just single-shot generation.
+"""
+
+from typing import Any, Dict, List, Optional
+from openenv.core.env_server.types import Action, Observation
+from pydantic import Field
+
+
+class SQLArenaAction(Action):
+    """
+    Action for SQLArenaEnv.
+
+    The agent has two action types:
+      - query_type="explore"  : Run an exploratory SQL query (no penalty, just observe)
+      - query_type="submit"   : Submit final answer query (ends episode, scored)
+
+    Example explore action:
+        SQLArenaAction(sql="SELECT * FROM orders LIMIT 5", query_type="explore")
+
+    Example submit action:
+        SQLArenaAction(sql="SELECT customer_id, SUM(amount) FROM orders GROUP BY customer_id", query_type="submit")
+    """
+
+    sql: str = Field(..., description="SQL query to execute against the database")
+    query_type: str = Field(
+        default="explore",
+        description="'explore' for investigative queries, 'submit' for final answer"
+    )
+
+
+class SQLArenaObservation(Observation):
+    """
+    Observation returned after each action in SQLArenaEnv.
+
+    Contains the result of the SQL execution, task context,
+    and feedback to guide the agent toward the correct answer.
+    """
+
+    # Task context (shown throughout episode)
+    task_id: str = Field(default="", description="Current task identifier")
+    difficulty: str = Field(default="easy", description="Task difficulty: easy/medium/hard/expert")
+    question: str = Field(default="", description="Natural language question to answer via SQL")
+    schema_info: str = Field(default="", description="Database schema description")
+
+    # Query result
+    query_result: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="Rows returned by the SQL query (up to 20 rows shown)"
+    )
+    query_error: Optional[str] = Field(
+        default=None,
+        description="SQL error message if query failed, null otherwise"
+    )
+    query_type: str = Field(default="explore", description="Type of query just executed")
+    rows_returned: int = Field(default=0, description="Number of rows returned")
+
+    # Episode progress
+    explore_steps_used: int = Field(default=0, description="Number of explore actions used so far")
+    explore_steps_remaining: int = Field(default=5, description="Explore actions remaining before forced submit")
+    submitted: bool = Field(default=False, description="Whether agent has submitted final answer")
+
+    # Feedback (only populated after submit)
+    is_correct: Optional[bool] = Field(default=None, description="Whether submitted answer is correct (after submit only)")
+    feedback: Optional[str] = Field(default=None, description="Feedback message after submission")
+    expected_row_count: Optional[int] = Field(default=None, description="How many rows correct answer has")

openenv.yaml

@@ -0,0 +1,6 @@
+spec_version: 1
+name: sql_arena_env
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000

pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-sql-arena-env"
+version = "1.0.0"
+description = "SQLArenaEnv — Multi-step SQL reasoning environment for OpenEnv. Agents explore a database, run investigative queries, and submit a final SQL answer. Tests true SQL understanding across 50 tasks in 4 difficulty tiers."
+requires-python = ">=3.10"
+dependencies = [
+    "openenv-core[core]>=0.2.2",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+]
+
+[project.scripts]
+server = "sql_arena_env.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["sql_arena_env", "sql_arena_env.server"]
+package-dir = { "sql_arena_env" = ".", "sql_arena_env.server" = "server" }

server/Dockerfile

@@ -0,0 +1,50 @@
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+ARG BUILD_MODE=standalone
+ARG ENV_NAME=sql_arena_env
+
+COPY . /app/env
+
+WORKDIR /app/env
+
+RUN if ! command -v uv >/dev/null 2>&1; then \
+    curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    mv /root/.local/bin/uv /usr/local/bin/uv && \
+    mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+    uv sync --frozen --no-install-project --no-editable; \
+    else \
+    uv sync --no-install-project --no-editable; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+    uv sync --frozen --no-editable; \
+    else \
+    uv sync --no-editable; \
+    fi
+
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+COPY --from=builder /app/env/.venv /app/.venv
+COPY --from=builder /app/env /app/env
+
+ENV PATH="/app/.venv/bin:$PATH"
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

server/__init__.py

@@ -0,0 +1 @@
+"""SQLArenaEnv server package."""

server/app.py

@@ -0,0 +1,51 @@
+"""
+FastAPI application for SQLArenaEnv.
+
+Exposes SQLArenaEnvironment over HTTP and WebSocket endpoints,
+fully compatible with OpenEnv EnvClient.
+
+Endpoints:
+    POST /reset    — start new episode (accepts task_id or difficulty in body)
+    POST /step     — submit SQL action
+    GET  /state    — get episode state
+    GET  /schema   — action/observation schema
+    GET  /tasks    — list all available tasks
+    GET  /health   — liveness check
+    WS   /ws       — persistent WebSocket session (preferred for training)
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:
+    raise ImportError(
+        "openenv-core is required. Install with: pip install openenv-core"
+    ) from e
+
+try:
+    from ..models import SQLArenaAction, SQLArenaObservation
+    from .sql_arena_environment import SQLArenaEnvironment
+except ModuleNotFoundError:
+    from models import SQLArenaAction, SQLArenaObservation
+    from server.sql_arena_environment import SQLArenaEnvironment
+
+
+app = create_app(
+    SQLArenaEnvironment,
+    SQLArenaAction,
+    SQLArenaObservation,
+    env_name="sql_arena_env",
+    max_concurrent_envs=4,  # support parallel grader runs
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    import uvicorn
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8000)
+    args = parser.parse_args()
+    main(port=args.port)

server/requirements.txt

@@ -0,0 +1,3 @@
+openenv-core[core]>=0.2.2
+fastapi>=0.115.0
+uvicorn>=0.24.0

server/sql_arena_environment.py

@@ -0,0 +1,382 @@
+"""
+SQLArenaEnv Core Environment.
+
+A multi-step SQL reasoning environment. The agent is given a natural language
+question about a database and must:
+  1. Run exploratory SQL queries to understand the schema and data (explore steps)
+  2. Submit a final SQL query that answers the question (submit step)
+
+The key novel mechanic: exploration queries are first-class actions.
+Agents that explore strategically score better than agents that blindly guess.
+
+Grading is done by comparing the agent's result set against the reference
+solution's result set — not string matching. This means equivalent SQL
+that produces the same answer is scored correctly.
+"""
+
+import sqlite3
+import json
+from typing import Any, Dict, List, Optional, Tuple
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+    from ..models import SQLArenaAction, SQLArenaObservation
+    from ..tasks import SQLTask, get_task, get_tasks_by_difficulty, ALL_TASKS
+except ImportError:
+    from models import SQLArenaAction, SQLArenaObservation
+    from tasks import SQLTask, get_task, get_tasks_by_difficulty, ALL_TASKS
+
+
+MAX_EXPLORE_STEPS = 5        # agent gets 5 free exploration queries
+MAX_RESULT_ROWS = 20         # cap result rows shown to agent
+MAX_QUERY_TIMEOUT_MS = 5000  # 5 second query timeout
+
+
+class SQLArenaEnvironment(Environment):
+    """
+    Multi-step SQL reasoning environment.
+
+    Each episode:
+      - A task is selected (by task_id or random from difficulty tier)
+      - Agent explores the database with up to MAX_EXPLORE_STEPS queries
+      - Agent submits a final answer query
+      - Reward is based on correctness of the final answer
+
+    Reward structure:
+      - Correct answer: 1.0
+      - Partially correct (right columns, wrong rows or ordering): 0.4
+      - Wrong answer: 0.0
+      - Exploration step: 0.0 (no reward, just information)
+      - Each explore step costs -0.02 (small penalty for over-exploration)
+      - Syntax error on submit: -0.1
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._conn: Optional[sqlite3.Connection] = None
+        self._current_task: Optional[SQLTask] = None
+        self._explore_steps_used: int = 0
+        self._submitted: bool = False
+        self._expected_result: Optional[List[Dict]] = None
+        self._task_id_override: Optional[str] = None
+        self._difficulty_override: Optional[str] = None
+
+    # ─────────────────────────────────────────
+    # OpenEnv required interface
+    # ─────────────────────────────────────────
+
+    def reset(self, task_id: Optional[str] = None, difficulty: Optional[str] = None) -> SQLArenaObservation:
+        """
+        Start a new episode with a fresh database and task.
+
+        Args:
+            task_id: specific task to load (e.g. 'easy_001')
+            difficulty: random task from this difficulty ('easy','medium','hard','expert')
+                        If neither given, picks a random task from any difficulty.
+        """
+        # Clean up old connection
+        if self._conn:
+            try:
+                self._conn.close()
+            except Exception:
+                pass
+
+        # Select task
+        if task_id:
+            task = get_task(task_id)
+            if not task:
+                task = list(ALL_TASKS.values())[0]
+        elif difficulty:
+            tasks = get_tasks_by_difficulty(difficulty)
+            import random
+            task = random.choice(tasks) if tasks else list(ALL_TASKS.values())[0]
+        else:
+            import random
+            task = random.choice(list(ALL_TASKS.values()))
+
+        self._current_task = task
+        self._explore_steps_used = 0
+        self._submitted = False
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+
+        # Create fresh in-memory SQLite database
+        self._conn = sqlite3.connect(":memory:", check_same_thread=False)
+        self._conn.row_factory = sqlite3.Row
+
+        # Load schema and seed data
+        try:
+            self._conn.executescript(task.schema_sql)
+            self._conn.executescript(task.seed_sql)
+            self._conn.commit()
+        except Exception as e:
+            return self._make_obs(
+                query_error=f"Environment setup error: {e}",
+                query_result=[],
+                query_type="reset",
+            )
+
+        # Pre-compute expected result
+        self._expected_result = self._run_query_safe(task.solution_sql)
+
+        return self._make_obs(
+            query_result=[],
+            query_type="reset",
+            query_error=None,
+        )
+
+    def step(self, action: SQLArenaAction) -> SQLArenaObservation:
+        """
+        Execute a SQL action.
+
+        action.query_type == "explore": run query, return results, no episode end
+        action.query_type == "submit":  run query, compare to answer, end episode
+        """
+        if self._current_task is None:
+            return self._make_obs(
+                query_error="Environment not initialized. Call reset() first.",
+                query_result=[],
+                query_type=action.query_type,
+                done=True,
+                reward=-1.0,
+            )
+
+        if self._submitted:
+            return self._make_obs(
+                query_error="Episode already ended. Call reset() to start a new episode.",
+                query_result=[],
+                query_type=action.query_type,
+                done=True,
+                reward=0.0,
+            )
+
+        self._state.step_count += 1
+        sql = action.sql.strip()
+        query_type = action.query_type.lower()
+
+        # ── EXPLORE action ──────────────────────────────
+        if query_type == "explore":
+            # Force submit if explore budget exhausted
+            if self._explore_steps_used >= MAX_EXPLORE_STEPS:
+                query_type = "submit"
+            else:
+                self._explore_steps_used += 1
+                result, error = self._execute_safe(sql)
+                # Small cost per explore step to discourage random fishing
+                reward = -0.02
+
+                return self._make_obs(
+                    query_result=result,
+                    query_error=error,
+                    query_type="explore",
+                    done=False,
+                    reward=reward,
+                )
+
+        # ── SUBMIT action ───────────────────────────────
+        self._submitted = True
+        result, error = self._execute_safe(sql)
+
+        if error:
+            # Syntax error on final submit
+            reward = -0.1
+            feedback = f"SQL error on submission: {error}. Correct your query."
+            return self._make_obs(
+                query_result=[],
+                query_error=error,
+                query_type="submit",
+                done=True,
+                reward=reward,
+                is_correct=False,
+                feedback=feedback,
+            )
+
+        # Grade the result
+        is_correct, partial, feedback = self._grade(result)
+        if is_correct:
+            reward = 1.0
+        elif partial:
+            reward = 0.4
+        else:
+            reward = 0.0
+
+        return self._make_obs(
+            query_result=result,
+            query_error=None,
+            query_type="submit",
+            done=True,
+            reward=reward,
+            is_correct=is_correct,
+            feedback=feedback,
+        )
+
+    @property
+    def state(self) -> State:
+        return self._state
+
+    # ─────────────────────────────────────────
+    # Internal helpers
+    # ─────────────────────────────────────────
+
+    def _execute_safe(self, sql: str) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+        """Execute SQL and return (rows, error). Caps rows at MAX_RESULT_ROWS."""
+        if not self._conn:
+            return [], "Database not initialized"
+        try:
+            # Block dangerous operations
+            sql_upper = sql.upper().strip()
+            dangerous = ["DROP ", "ALTER ", "TRUNCATE ", "PRAGMA ", "ATTACH ", "DETACH "]
+            if any(sql_upper.startswith(d) for d in dangerous):
+                return [], "Operation not permitted in this environment"
+
+            cursor = self._conn.execute(sql)
+            rows = cursor.fetchmany(MAX_RESULT_ROWS + 1)
+            truncated = len(rows) > MAX_RESULT_ROWS
+            rows = rows[:MAX_RESULT_ROWS]
+            result = [dict(row) for row in rows]
+            if truncated:
+                result.append({"__info__": f"Results truncated to {MAX_RESULT_ROWS} rows"})
+            return result, None
+        except sqlite3.Error as e:
+            return [], str(e)
+        except Exception as e:
+            return [], f"Unexpected error: {e}"
+
+    def _run_query_safe(self, sql: str) -> Optional[List[Dict[str, Any]]]:
+        """Run reference solution, return all rows."""
+        if not self._conn:
+            return None
+        try:
+            cursor = self._conn.execute(sql)
+            rows = cursor.fetchall()
+            return [dict(row) for row in rows]
+        except Exception:
+            return None
+
+    def _grade(self, agent_result: List[Dict]) -> Tuple[bool, bool, str]:
+        """
+        Grade agent's result against expected result.
+        Returns (is_correct, is_partial, feedback_message).
+
+        Grading logic:
+          - Full credit: same rows, same values (order-insensitive for most tasks)
+          - Partial credit: same column names but different rows, OR right count wrong values
+          - No credit: completely wrong
+        """
+        expected = self._expected_result
+        if expected is None:
+            return False, False, "Could not compute expected result. Contact organizers."
+
+        # Normalize both result sets for comparison
+        def normalize(rows: List[Dict]) -> List[str]:
+            """Convert rows to sorted list of JSON strings for comparison."""
+            normalized = []
+            for row in rows:
+                # Remove __info__ markers
+                clean = {k: v for k, v in row.items() if not k.startswith("__")}
+                # Round floats to 2 decimal places to avoid floating point issues
+                rounded = {}
+                for k, v in clean.items():
+                    if isinstance(v, float):
+                        rounded[k] = round(v, 2)
+                    else:
+                        rounded[k] = v
+                normalized.append(json.dumps(rounded, sort_keys=True))
+            return sorted(normalized)
+
+        # Filter out info rows from agent result
+        agent_clean = [r for r in agent_result if not any(k.startswith("__") for k in r.keys())]
+
+        expected_norm = normalize(expected)
+        agent_norm = normalize(agent_clean)
+
+        # Full correct
+        if expected_norm == agent_norm:
+            return True, False, f"✓ Correct! Your query returned the exact expected result ({len(expected)} rows)."
+
+        # Check partial credit conditions
+        # Same number of rows?
+        same_count = len(agent_clean) == len(expected)
+
+        # Same column names?
+        exp_cols = set(expected[0].keys()) if expected else set()
+        agent_cols = set(agent_clean[0].keys()) if agent_clean else set()
+        same_cols = exp_cols == agent_cols
+
+        # How many rows match?
+        matching_rows = len(set(expected_norm) & set(agent_norm))
+        match_pct = matching_rows / max(len(expected_norm), 1) * 100
+
+        if same_cols and same_count and match_pct >= 50:
+            return False, True, (
+                f"Partial credit. Correct columns, {same_count} rows, "
+                f"but {matching_rows}/{len(expected)} rows match exactly. "
+                f"Check your WHERE conditions or aggregation."
+            )
+
+        if same_cols and not same_count:
+            return False, True, (
+                f"Partial credit. Correct columns but wrong row count: "
+                f"got {len(agent_clean)}, expected {len(expected)}. "
+                f"Check your filters."
+            ) if match_pct >= 30 else (False, False, (
+                f"Wrong answer. Got {len(agent_clean)} rows, expected {len(expected)}. "
+                f"Expected columns: {sorted(exp_cols)}."
+            ))
+
+        feedback = (
+            f"Wrong answer. Expected {len(expected)} rows with columns {sorted(exp_cols)}. "
+            f"You returned {len(agent_clean)} rows"
+            + (f" with columns {sorted(agent_cols)}." if agent_clean else " (empty result).")
+        )
+        return False, False, feedback
+
+    def _make_obs(
+        self,
+        query_result: List[Dict],
+        query_type: str,
+        query_error: Optional[str] = None,
+        done: bool = False,
+        reward: float = 0.0,
+        is_correct: Optional[bool] = None,
+        feedback: Optional[str] = None,
+    ) -> SQLArenaObservation:
+        """Construct a full observation."""
+        task = self._current_task
+        expected_count = len(self._expected_result) if self._expected_result else None
+
+        return SQLArenaObservation(
+            # Task context
+            task_id=task.task_id if task else "",
+            difficulty=task.difficulty if task else "easy",
+            question=task.question if task else "",
+            schema_info=task.schema_description if task else "",
+
+            # Query result
+            query_result=query_result,
+            query_error=query_error,
+            query_type=query_type,
+            rows_returned=len([r for r in query_result if not any(k.startswith("__") for k in r.keys())]),
+
+            # Episode progress
+            explore_steps_used=self._explore_steps_used,
+            explore_steps_remaining=max(0, MAX_EXPLORE_STEPS - self._explore_steps_used),
+            submitted=self._submitted,
+
+            # Feedback after submit
+            is_correct=is_correct,
+            feedback=feedback,
+            expected_row_count=expected_count,
+
+            # OpenEnv base fields
+            done=done,
+            reward=reward,
+            metadata={
+                "episode_id": self._state.episode_id,
+                "step_count": self._state.step_count,
+            },
+        )

tasks.py

@@ -0,0 +1,1348 @@
+"""
+SQLArenaEnv Task Library.
+
+50 hand-crafted SQL tasks across 4 difficulty tiers.
+Each task includes:
+  - A SQLite schema (CREATE TABLE statements)
+  - Seed data (INSERT statements)
+  - A natural language question
+  - A validator function that checks agent's result against correct answer
+
+Difficulty tiers:
+  EASY   (10 tasks) — single table, basic SELECT/WHERE/ORDER BY
+  MEDIUM (15 tasks) — JOINs, GROUP BY, HAVING, subqueries
+  HARD   (15 tasks) — window functions, CTEs, multi-join, aggregation
+  EXPERT (10 tasks) — correlated subqueries, recursive CTEs, complex analytics
+"""
+
+from typing import Any, Callable, Dict, List, Optional
+from dataclasses import dataclass, field
+
+
+@dataclass
+class SQLTask:
+    task_id: str
+    difficulty: str  # easy / medium / hard / expert
+    question: str
+    schema_sql: str          # CREATE TABLE statements
+    seed_sql: str            # INSERT statements
+    solution_sql: str        # Reference solution (for generating expected answer)
+    schema_description: str  # Human-readable schema hint for agent
+    validator: Optional[Callable] = None  # custom validator, None = use default row comparison
+
+
+# ─────────────────────────────────────────────
+# EASY TASKS — single table operations
+# ─────────────────────────────────────────────
+
+EASY_TASKS = [
+    SQLTask(
+        task_id="easy_001",
+        difficulty="easy",
+        question="List all employees whose salary is greater than 70000, ordered by salary descending.",
+        schema_sql="""
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department TEXT,
+    salary REAL,
+    hire_date TEXT
+);""",
+        seed_sql="""
+INSERT INTO employees VALUES
+(1,'Alice','Engineering',95000,'2020-01-15'),
+(2,'Bob','Marketing',55000,'2019-03-20'),
+(3,'Charlie','Engineering',82000,'2021-06-01'),
+(4,'Diana','HR',47000,'2018-11-10'),
+(5,'Eve','Engineering',110000,'2017-08-25'),
+(6,'Frank','Marketing',72000,'2022-02-14'),
+(7,'Grace','HR',68000,'2020-09-30'),
+(8,'Hank','Engineering',78000,'2019-12-05');""",
+        solution_sql="SELECT id, name, department, salary FROM employees WHERE salary > 70000 ORDER BY salary DESC;",
+        schema_description="employees(id, name, department, salary, hire_date)",
+    ),
+
+    SQLTask(
+        task_id="easy_002",
+        difficulty="easy",
+        question="How many products are there in each category? Show category and count, ordered by count descending.",
+        schema_sql="""
+CREATE TABLE products (
+    product_id INTEGER PRIMARY KEY,
+    name TEXT,
+    category TEXT,
+    price REAL,
+    stock INTEGER
+);""",
+        seed_sql="""
+INSERT INTO products VALUES
+(1,'Laptop','Electronics',999.99,50),
+(2,'Phone','Electronics',699.99,120),
+(3,'Desk','Furniture',249.99,30),
+(4,'Chair','Furniture',189.99,75),
+(5,'Tablet','Electronics',449.99,60),
+(6,'Bookshelf','Furniture',129.99,40),
+(7,'Headphones','Electronics',149.99,200),
+(8,'Lamp','Furniture',59.99,90),
+(9,'Monitor','Electronics',399.99,45),
+(10,'Keyboard','Electronics',79.99,150);""",
+        solution_sql="SELECT category, COUNT(*) as count FROM products GROUP BY category ORDER BY count DESC;",
+        schema_description="products(product_id, name, category, price, stock)",
+    ),
+
+    SQLTask(
+        task_id="easy_003",
+        difficulty="easy",
+        question="Find the top 5 most expensive products. Show name and price.",
+        schema_sql="""
+CREATE TABLE products (
+    product_id INTEGER PRIMARY KEY,
+    name TEXT,
+    category TEXT,
+    price REAL,
+    stock INTEGER
+);""",
+        seed_sql="""
+INSERT INTO products VALUES
+(1,'Laptop','Electronics',999.99,50),
+(2,'Phone','Electronics',699.99,120),
+(3,'Desk','Furniture',249.99,30),
+(4,'Chair','Furniture',189.99,75),
+(5,'Tablet','Electronics',449.99,60),
+(6,'Bookshelf','Furniture',129.99,40),
+(7,'Headphones','Electronics',149.99,200),
+(8,'Lamp','Furniture',59.99,90),
+(9,'Monitor','Electronics',399.99,45),
+(10,'Keyboard','Electronics',79.99,150);""",
+        solution_sql="SELECT name, price FROM products ORDER BY price DESC LIMIT 5;",
+        schema_description="products(product_id, name, category, price, stock)",
+    ),
+
+    SQLTask(
+        task_id="easy_004",
+        difficulty="easy",
+        question="Find all students who scored above 85 in Mathematics. Show name and score.",
+        schema_sql="""
+CREATE TABLE students (
+    student_id INTEGER PRIMARY KEY,
+    name TEXT,
+    subject TEXT,
+    score INTEGER,
+    grade TEXT
+);""",
+        seed_sql="""
+INSERT INTO students VALUES
+(1,'Arjun','Mathematics',92,'A'),
+(2,'Priya','Mathematics',78,'B'),
+(3,'Rahul','Science',88,'A'),
+(4,'Sneha','Mathematics',95,'A'),
+(5,'Vikram','Science',72,'C'),
+(6,'Ananya','Mathematics',65,'D'),
+(7,'Karan','Mathematics',89,'A'),
+(8,'Divya','Science',91,'A'),
+(9,'Rohit','Mathematics',55,'F'),
+(10,'Meera','Mathematics',87,'A');""",
+        solution_sql="SELECT name, score FROM students WHERE subject='Mathematics' AND score > 85 ORDER BY score DESC;",
+        schema_description="students(student_id, name, subject, score, grade)",
+    ),
+
+    SQLTask(
+        task_id="easy_005",
+        difficulty="easy",
+        question="What is the total revenue (price * quantity) from all orders? Return a single number called total_revenue.",
+        schema_sql="""
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    customer TEXT,
+    product TEXT,
+    price REAL,
+    quantity INTEGER,
+    order_date TEXT
+);""",
+        seed_sql="""
+INSERT INTO orders VALUES
+(1,'Alice','Laptop',999.99,2,'2024-01-10'),
+(2,'Bob','Phone',699.99,1,'2024-01-11'),
+(3,'Charlie','Tablet',449.99,3,'2024-01-12'),
+(4,'Diana','Headphones',149.99,5,'2024-01-13'),
+(5,'Eve','Monitor',399.99,1,'2024-01-14'),
+(6,'Frank','Keyboard',79.99,4,'2024-01-15'),
+(7,'Grace','Laptop',999.99,1,'2024-01-16'),
+(8,'Hank','Phone',699.99,2,'2024-01-17');""",
+        solution_sql="SELECT ROUND(SUM(price * quantity), 2) as total_revenue FROM orders;",
+        schema_description="orders(order_id, customer, product, price, quantity, order_date)",
+    ),
+
+    SQLTask(
+        task_id="easy_006",
+        difficulty="easy",
+        question="List all cities that have more than 1 million population, ordered alphabetically.",
+        schema_sql="""
+CREATE TABLE cities (
+    city_id INTEGER PRIMARY KEY,
+    name TEXT,
+    country TEXT,
+    population INTEGER,
+    area_km2 REAL
+);""",
+        seed_sql="""
+INSERT INTO cities VALUES
+(1,'Mumbai','India',20667656,603.4),
+(2,'Delhi','India',32941000,1484.0),
+(3,'Kolkata','India',14850000,206.1),
+(4,'Chennai','India',10971000,426.0),
+(5,'Pune','India',3124000,331.3),
+(6,'Hyderabad','India',10534000,650.0),
+(7,'Ahmedabad','India',8450000,475.0),
+(8,'Jaipur','India',3766000,467.0),
+(9,'Surat','India',6936000,326.5),
+(10,'Visakhapatnam','India',2035000,681.96);""",
+        solution_sql="SELECT name, population FROM cities WHERE population > 1000000 ORDER BY name ASC;",
+        schema_description="cities(city_id, name, country, population, area_km2)",
+    ),
+
+    SQLTask(
+        task_id="easy_007",
+        difficulty="easy",
+        question="Find the average salary by department. Show department and avg_salary rounded to 2 decimals.",
+        schema_sql="""
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department TEXT,
+    salary REAL,
+    hire_date TEXT
+);""",
+        seed_sql="""
+INSERT INTO employees VALUES
+(1,'Alice','Engineering',95000,'2020-01-15'),
+(2,'Bob','Marketing',55000,'2019-03-20'),
+(3,'Charlie','Engineering',82000,'2021-06-01'),
+(4,'Diana','HR',47000,'2018-11-10'),
+(5,'Eve','Engineering',110000,'2017-08-25'),
+(6,'Frank','Marketing',72000,'2022-02-14'),
+(7,'Grace','HR',68000,'2020-09-30'),
+(8,'Hank','Engineering',78000,'2019-12-05'),
+(9,'Ivy','Marketing',61000,'2021-03-15'),
+(10,'Jack','HR',52000,'2022-07-01');""",
+        solution_sql="SELECT department, ROUND(AVG(salary),2) as avg_salary FROM employees GROUP BY department ORDER BY avg_salary DESC;",
+        schema_description="employees(id, name, department, salary, hire_date)",
+    ),
+
+    SQLTask(
+        task_id="easy_008",
+        difficulty="easy",
+        question="Find all books published after 2010 with rating above 4.0. Show title, author, and rating.",
+        schema_sql="""
+CREATE TABLE books (
+    book_id INTEGER PRIMARY KEY,
+    title TEXT,
+    author TEXT,
+    year INTEGER,
+    rating REAL,
+    genre TEXT
+);""",
+        seed_sql="""
+INSERT INTO books VALUES
+(1,'The Pragmatic Programmer','Hunt & Thomas',1999,4.8,'Tech'),
+(2,'Clean Code','Robert Martin',2008,4.6,'Tech'),
+(3,'Thinking Fast and Slow','Daniel Kahneman',2011,4.5,'Psychology'),
+(4,'Atomic Habits','James Clear',2018,4.8,'Self-Help'),
+(5,'Deep Work','Cal Newport',2016,4.4,'Self-Help'),
+(6,'The Lean Startup','Eric Ries',2011,4.2,'Business'),
+(7,'Zero to One','Peter Thiel',2014,4.3,'Business'),
+(8,'Dune','Frank Herbert',1965,4.7,'Fiction'),
+(9,'Project Hail Mary','Andy Weir',2021,4.9,'Fiction'),
+(10,'The Midnight Library','Matt Haig',2020,4.2,'Fiction');""",
+        solution_sql="SELECT title, author, rating FROM books WHERE year > 2010 AND rating > 4.0 ORDER BY rating DESC;",
+        schema_description="books(book_id, title, author, year, rating, genre)",
+    ),
+
+    SQLTask(
+        task_id="easy_009",
+        difficulty="easy",
+        question="Count how many employees were hired each year. Show hire_year and count, ordered by year.",
+        schema_sql="""
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department TEXT,
+    salary REAL,
+    hire_date TEXT
+);""",
+        seed_sql="""
+INSERT INTO employees VALUES
+(1,'Alice','Engineering',95000,'2020-01-15'),
+(2,'Bob','Marketing',55000,'2019-03-20'),
+(3,'Charlie','Engineering',82000,'2021-06-01'),
+(4,'Diana','HR',47000,'2018-11-10'),
+(5,'Eve','Engineering',110000,'2017-08-25'),
+(6,'Frank','Marketing',72000,'2022-02-14'),
+(7,'Grace','HR',68000,'2020-09-30'),
+(8,'Hank','Engineering',78000,'2019-12-05'),
+(9,'Ivy','Marketing',61000,'2021-03-15'),
+(10,'Jack','HR',52000,'2022-07-01');""",
+        solution_sql="SELECT SUBSTR(hire_date,1,4) as hire_year, COUNT(*) as count FROM employees GROUP BY hire_year ORDER BY hire_year;",
+        schema_description="employees(id, name, department, salary, hire_date — format YYYY-MM-DD)",
+    ),
+
+    SQLTask(
+        task_id="easy_010",
+        difficulty="easy",
+        question="Find the minimum and maximum price for each product category.",
+        schema_sql="""
+CREATE TABLE products (
+    product_id INTEGER PRIMARY KEY,
+    name TEXT,
+    category TEXT,
+    price REAL,
+    stock INTEGER
+);""",
+        seed_sql="""
+INSERT INTO products VALUES
+(1,'Laptop','Electronics',999.99,50),
+(2,'Phone','Electronics',699.99,120),
+(3,'Desk','Furniture',249.99,30),
+(4,'Chair','Furniture',189.99,75),
+(5,'Tablet','Electronics',449.99,60),
+(6,'Bookshelf','Furniture',129.99,40),
+(7,'Headphones','Electronics',149.99,200),
+(8,'Lamp','Furniture',59.99,90),
+(9,'Monitor','Electronics',399.99,45),
+(10,'Keyboard','Electronics',79.99,150);""",
+        solution_sql="SELECT category, MIN(price) as min_price, MAX(price) as max_price FROM products GROUP BY category ORDER BY category;",
+        schema_description="products(product_id, name, category, price, stock)",
+    ),
+]
+
+
+# ─────────────────────────────────────────────
+# MEDIUM TASKS — JOINs, GROUP BY, HAVING
+# ─────────────────────────────────────────────
+
+MEDIUM_SCHEMA_ECOMMERCE = """
+CREATE TABLE customers (
+    customer_id INTEGER PRIMARY KEY,
+    name TEXT,
+    email TEXT,
+    city TEXT,
+    joined_date TEXT
+);
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    customer_id INTEGER,
+    order_date TEXT,
+    status TEXT,
+    FOREIGN KEY (customer_id) REFERENCES customers(customer_id)
+);
+CREATE TABLE order_items (
+    item_id INTEGER PRIMARY KEY,
+    order_id INTEGER,
+    product_name TEXT,
+    quantity INTEGER,
+    unit_price REAL,
+    FOREIGN KEY (order_id) REFERENCES orders(order_id)
+);"""
+
+MEDIUM_SEED_ECOMMERCE = """
+INSERT INTO customers VALUES
+(1,'Arjun Sharma','arjun@email.com','Mumbai','2022-01-15'),
+(2,'Priya Patel','priya@email.com','Delhi','2022-03-20'),
+(3,'Rahul Gupta','rahul@email.com','Bangalore','2021-11-05'),
+(4,'Sneha Reddy','sneha@email.com','Hyderabad','2023-02-10'),
+(5,'Vikram Singh','vikram@email.com','Chennai','2021-08-25'),
+(6,'Ananya Joshi','ananya@email.com','Pune','2022-07-14'),
+(7,'Karan Mehta','karan@email.com','Mumbai','2023-01-30'),
+(8,'Divya Nair','divya@email.com','Bangalore','2022-05-18');
+
+INSERT INTO orders VALUES
+(101,1,'2024-01-10','completed'),
+(102,1,'2024-02-15','completed'),
+(103,2,'2024-01-20','completed'),
+(104,3,'2024-01-25','pending'),
+(105,3,'2024-02-28','completed'),
+(106,4,'2024-03-05','completed'),
+(107,5,'2024-01-08','cancelled'),
+(108,5,'2024-03-12','completed'),
+(109,6,'2024-02-20','completed'),
+(110,7,'2024-03-01','pending'),
+(111,8,'2024-01-30','completed'),
+(112,8,'2024-02-25','completed');
+
+INSERT INTO order_items VALUES
+(1,101,'Laptop',1,999.99),
+(2,101,'Mouse',2,29.99),
+(3,102,'Phone',1,699.99),
+(4,103,'Tablet',2,449.99),
+(5,104,'Headphones',3,149.99),
+(6,105,'Monitor',1,399.99),
+(7,106,'Keyboard',2,79.99),
+(8,107,'Laptop',1,999.99),
+(9,108,'Phone',1,699.99),
+(10,109,'Desk',1,249.99),
+(11,110,'Chair',2,189.99),
+(12,111,'Laptop',1,999.99),
+(13,111,'Tablet',1,449.99),
+(14,112,'Headphones',2,149.99);"""
+
+MEDIUM_SCHEMA_HR = """
+CREATE TABLE departments (
+    dept_id INTEGER PRIMARY KEY,
+    dept_name TEXT,
+    manager_id INTEGER,
+    budget REAL
+);
+CREATE TABLE employees (
+    emp_id INTEGER PRIMARY KEY,
+    name TEXT,
+    dept_id INTEGER,
+    salary REAL,
+    hire_date TEXT,
+    manager_id INTEGER,
+    FOREIGN KEY (dept_id) REFERENCES departments(dept_id)
+);"""
+
+MEDIUM_SEED_HR = """
+INSERT INTO departments VALUES
+(1,'Engineering',5,5000000),
+(2,'Marketing',8,2000000),
+(3,'HR',12,1500000),
+(4,'Finance',15,3000000),
+(5,'Operations',18,2500000);
+
+INSERT INTO employees VALUES
+(1,'Alice Chen',1,95000,'2020-01-15',5),
+(2,'Bob Kumar',1,82000,'2021-06-01',5),
+(3,'Charlie Roy',1,78000,'2019-12-05',5),
+(4,'Diana Shah',2,72000,'2022-02-14',8),
+(5,'Eve Patel',1,110000,'2017-08-25',NULL),
+(6,'Frank Joshi',2,55000,'2019-03-20',8),
+(7,'Grace Singh',3,68000,'2020-09-30',12),
+(8,'Hank Mehta',2,88000,'2018-06-15',NULL),
+(9,'Ivy Reddy',3,52000,'2022-07-01',12),
+(10,'Jack Nair',3,47000,'2018-11-10',12),
+(11,'Kara Gupta',4,91000,'2019-04-22',15),
+(12,'Leo Sharma',3,75000,'2016-03-10',NULL),
+(13,'Mia Verma',4,83000,'2020-08-14',15),
+(14,'Nia Bose',5,69000,'2021-01-20',18),
+(15,'Omar Das',4,102000,'2015-11-30',NULL),
+(16,'Pia Iyer',5,61000,'2022-03-25',18),
+(17,'Quinn Rao',1,74000,'2021-09-08',5),
+(18,'Rita Pillai',5,95000,'2016-07-12',NULL);"""
+
+MEDIUM_TASKS = [
+    SQLTask(
+        task_id="medium_001",
+        difficulty="medium",
+        question="Find customers who have placed more than 1 completed order. Show customer name and order count.",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT c.name, COUNT(o.order_id) as order_count
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+WHERE o.status = 'completed'
+GROUP BY c.customer_id, c.name
+HAVING COUNT(o.order_id) > 1
+ORDER BY order_count DESC;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_002",
+        difficulty="medium",
+        question="Calculate total spending per customer (only completed orders). Show customer name and total_spent, ordered by total_spent descending.",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT c.name, ROUND(SUM(oi.quantity * oi.unit_price), 2) as total_spent
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+JOIN order_items oi ON o.order_id = oi.order_id
+WHERE o.status = 'completed'
+GROUP BY c.customer_id, c.name
+ORDER BY total_spent DESC;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_003",
+        difficulty="medium",
+        question="Find the most popular product (by total quantity sold across all completed orders).",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT oi.product_name, SUM(oi.quantity) as total_sold
+FROM order_items oi
+JOIN orders o ON oi.order_id = o.order_id
+WHERE o.status = 'completed'
+GROUP BY oi.product_name
+ORDER BY total_sold DESC
+LIMIT 1;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_004",
+        difficulty="medium",
+        question="Find departments where average salary exceeds 80000. Show dept_name and avg_salary.",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT d.dept_name, ROUND(AVG(e.salary), 2) as avg_salary
+FROM departments d
+JOIN employees e ON d.dept_id = e.dept_id
+GROUP BY d.dept_id, d.dept_name
+HAVING AVG(e.salary) > 80000
+ORDER BY avg_salary DESC;""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id)",
+    ),
+
+    SQLTask(
+        task_id="medium_005",
+        difficulty="medium",
+        question="For each department, find the highest-paid employee. Show dept_name, employee name, and salary.",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT d.dept_name, e.name, e.salary
+FROM employees e
+JOIN departments d ON e.dept_id = d.dept_id
+WHERE e.salary = (
+    SELECT MAX(e2.salary) FROM employees e2 WHERE e2.dept_id = e.dept_id
+)
+ORDER BY e.salary DESC;""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id)",
+    ),
+
+    SQLTask(
+        task_id="medium_006",
+        difficulty="medium",
+        question="List customers from Mumbai or Bangalore who have at least one order. Show name and city.",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT DISTINCT c.name, c.city
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+WHERE c.city IN ('Mumbai','Bangalore')
+ORDER BY c.city, c.name;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_007",
+        difficulty="medium",
+        question="Find employees who earn more than their department's average salary. Show name, salary, and dept_name.",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT e.name, e.salary, d.dept_name
+FROM employees e
+JOIN departments d ON e.dept_id = d.dept_id
+WHERE e.salary > (
+    SELECT AVG(e2.salary) FROM employees e2 WHERE e2.dept_id = e.dept_id
+)
+ORDER BY d.dept_name, e.salary DESC;""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id)",
+    ),
+
+    SQLTask(
+        task_id="medium_008",
+        difficulty="medium",
+        question="Calculate the cancellation rate per customer (cancelled orders / total orders). Show only customers with cancellation rate > 0. Show name and cancellation_rate rounded to 2 decimals.",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT c.name,
+       ROUND(CAST(SUM(CASE WHEN o.status='cancelled' THEN 1 ELSE 0 END) AS REAL) / COUNT(o.order_id), 2) as cancellation_rate
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+GROUP BY c.customer_id, c.name
+HAVING SUM(CASE WHEN o.status='cancelled' THEN 1 ELSE 0 END) > 0
+ORDER BY cancellation_rate DESC;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_009",
+        difficulty="medium",
+        question="Find total salary cost per department and what percentage of total company salary each dept represents. Show dept_name, total_salary, salary_pct rounded to 1 decimal.",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT d.dept_name,
+       SUM(e.salary) as total_salary,
+       ROUND(SUM(e.salary) * 100.0 / (SELECT SUM(salary) FROM employees), 1) as salary_pct
+FROM employees e
+JOIN departments d ON e.dept_id = d.dept_id
+GROUP BY d.dept_id, d.dept_name
+ORDER BY total_salary DESC;""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id)",
+    ),
+
+    SQLTask(
+        task_id="medium_010",
+        difficulty="medium",
+        question="Find orders placed in January 2024. Show customer name, order_date, and order total (sum of quantity*unit_price).",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT c.name, o.order_date, ROUND(SUM(oi.quantity * oi.unit_price),2) as order_total
+FROM orders o
+JOIN customers c ON o.customer_id = c.customer_id
+JOIN order_items oi ON o.order_id = oi.order_id
+WHERE o.order_date LIKE '2024-01%'
+GROUP BY o.order_id, c.name, o.order_date
+ORDER BY o.order_date;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_011",
+        difficulty="medium",
+        question="Which product has generated the most revenue across all completed orders? Show product_name and total_revenue.",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT oi.product_name, ROUND(SUM(oi.quantity * oi.unit_price),2) as total_revenue
+FROM order_items oi
+JOIN orders o ON oi.order_id = o.order_id
+WHERE o.status = 'completed'
+GROUP BY oi.product_name
+ORDER BY total_revenue DESC
+LIMIT 1;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+
+    SQLTask(
+        task_id="medium_012",
+        difficulty="medium",
+        question="Find employees who have been at the company for more than 4 years (hired before 2021). Show name, hire_date, and dept_name.",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT e.name, e.hire_date, d.dept_name
+FROM employees e
+JOIN departments d ON e.dept_id = d.dept_id
+WHERE e.hire_date < '2021-01-01'
+ORDER BY e.hire_date;""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id)",
+    ),
+
+    SQLTask(
+        task_id="medium_013",
+        difficulty="medium",
+        question="Find the second highest salary in the company (single value, call it second_highest_salary).",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT MAX(salary) as second_highest_salary
+FROM employees
+WHERE salary < (SELECT MAX(salary) FROM employees);""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id)",
+    ),
+
+    SQLTask(
+        task_id="medium_014",
+        difficulty="medium",
+        question="Count the number of employees per manager (only show managers with 2+ direct reports). Show manager name and report_count.",
+        schema_sql=MEDIUM_SCHEMA_HR,
+        seed_sql=MEDIUM_SEED_HR,
+        solution_sql="""
+SELECT m.name as manager_name, COUNT(e.emp_id) as report_count
+FROM employees e
+JOIN employees m ON e.manager_id = m.emp_id
+GROUP BY m.emp_id, m.name
+HAVING COUNT(e.emp_id) >= 2
+ORDER BY report_count DESC;""",
+        schema_description="departments(dept_id, dept_name, manager_id, budget), employees(emp_id, name, dept_id, salary, hire_date, manager_id) — manager_id refers to emp_id of the manager",
+    ),
+
+    SQLTask(
+        task_id="medium_015",
+        difficulty="medium",
+        question="Find customers who have never placed an order. Show their name and email.",
+        schema_sql=MEDIUM_SCHEMA_ECOMMERCE,
+        seed_sql=MEDIUM_SEED_ECOMMERCE,
+        solution_sql="""
+SELECT c.name, c.email
+FROM customers c
+LEFT JOIN orders o ON c.customer_id = o.customer_id
+WHERE o.order_id IS NULL
+ORDER BY c.name;""",
+        schema_description="customers(customer_id, name, email, city, joined_date), orders(order_id, customer_id, order_date, status), order_items(item_id, order_id, product_name, quantity, unit_price)",
+    ),
+]
+
+
+# ─────────────────────────────────────────────
+# HARD TASKS — CTEs, Window Functions
+# ─────────────────────────────────────────────
+
+HARD_SCHEMA = """
+CREATE TABLE customers (
+    customer_id INTEGER PRIMARY KEY,
+    name TEXT, city TEXT, segment TEXT
+);
+CREATE TABLE orders (
+    order_id INTEGER PRIMARY KEY,
+    customer_id INTEGER,
+    order_date TEXT,
+    ship_date TEXT,
+    category TEXT,
+    amount REAL
+);
+CREATE TABLE returns (
+    return_id INTEGER PRIMARY KEY,
+    order_id INTEGER,
+    return_date TEXT,
+    reason TEXT
+);"""
+
+HARD_SEED = """
+INSERT INTO customers VALUES
+(1,'Arjun','Mumbai','Consumer'),
+(2,'Priya','Delhi','Corporate'),
+(3,'Rahul','Bangalore','Consumer'),
+(4,'Sneha','Hyderabad','Corporate'),
+(5,'Vikram','Chennai','Home Office'),
+(6,'Ananya','Pune','Consumer'),
+(7,'Karan','Mumbai','Corporate'),
+(8,'Divya','Bangalore','Consumer');
+
+INSERT INTO orders VALUES
+(1001,1,'2023-01-10','2023-01-15','Technology',2500.00),
+(1002,1,'2023-03-20','2023-03-25','Furniture',1200.00),
+(1003,1,'2023-06-15','2023-06-20','Technology',3800.00),
+(1004,2,'2023-02-05','2023-02-10','Office Supplies',450.00),
+(1005,2,'2023-04-18','2023-04-23','Technology',5200.00),
+(1006,2,'2023-07-30','2023-08-04','Furniture',890.00),
+(1007,3,'2023-01-25','2023-01-30','Technology',1100.00),
+(1008,3,'2023-05-12','2023-05-17','Office Supplies',320.00),
+(1009,4,'2023-03-08','2023-03-13','Furniture',2200.00),
+(1010,4,'2023-08-22','2023-08-27','Technology',4100.00),
+(1011,5,'2023-02-14','2023-02-19','Office Supplies',180.00),
+(1012,5,'2023-06-28','2023-07-03','Technology',2900.00),
+(1013,6,'2023-01-05','2023-01-10','Furniture',750.00),
+(1014,6,'2023-04-25','2023-04-30','Office Supplies',290.00),
+(1015,7,'2023-03-15','2023-03-20','Technology',6800.00),
+(1016,7,'2023-07-10','2023-07-15','Furniture',1500.00),
+(1017,8,'2023-02-28','2023-03-05','Technology',3200.00),
+(1018,8,'2023-05-20','2023-05-25','Office Supplies',410.00);
+
+INSERT INTO returns VALUES
+(1,1002,'2023-04-01','Defective'),
+(2,1004,'2023-02-20','Wrong item'),
+(3,1008,'2023-05-25','Not needed'),
+(4,1011,'2023-02-25','Defective'),
+(5,1013,'2023-01-20','Wrong item');"""
+
+HARD_TASKS = [
+    SQLTask(
+        task_id="hard_001",
+        difficulty="hard",
+        question="Using a window function, rank customers by their total order amount within each city segment. Show customer name, city, total_amount, and their rank within city.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT c.name, c.city,
+       SUM(o.amount) as total_amount,
+       RANK() OVER (PARTITION BY c.city ORDER BY SUM(o.amount) DESC) as city_rank
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+GROUP BY c.customer_id, c.name, c.city
+ORDER BY c.city, city_rank;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_002",
+        difficulty="hard",
+        question="Using a CTE, find the running total of order amounts for each customer ordered by order_date. Show customer name, order_date, amount, and running_total.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+WITH customer_orders AS (
+    SELECT c.name, o.order_date, o.amount,
+           SUM(o.amount) OVER (PARTITION BY c.customer_id ORDER BY o.order_date) as running_total
+    FROM customers c
+    JOIN orders o ON c.customer_id = o.customer_id
+)
+SELECT name, order_date, amount, running_total
+FROM customer_orders
+ORDER BY name, order_date;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_003",
+        difficulty="hard",
+        question="Find the return rate by category (returned orders / total orders). Show category and return_rate as a percentage rounded to 1 decimal.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT o.category,
+       ROUND(COUNT(r.return_id) * 100.0 / COUNT(o.order_id), 1) as return_rate
+FROM orders o
+LEFT JOIN returns r ON o.order_id = r.order_id
+GROUP BY o.category
+ORDER BY return_rate DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_004",
+        difficulty="hard",
+        question="For each customer, find the month with their highest single order amount. Show customer name, best_month (YYYY-MM format), and max_amount.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+WITH monthly_max AS (
+    SELECT c.name, SUBSTR(o.order_date,1,7) as month, MAX(o.amount) as max_amount,
+           RANK() OVER (PARTITION BY c.customer_id ORDER BY MAX(o.amount) DESC) as rnk
+    FROM customers c
+    JOIN orders o ON c.customer_id = o.customer_id
+    GROUP BY c.customer_id, c.name, month
+)
+SELECT name, month as best_month, max_amount
+FROM monthly_max
+WHERE rnk = 1
+ORDER BY max_amount DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_005",
+        difficulty="hard",
+        question="Calculate average shipping time (days between order_date and ship_date) per category. Show category and avg_ship_days rounded to 1 decimal.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT category,
+       ROUND(AVG(julianday(ship_date) - julianday(order_date)), 1) as avg_ship_days
+FROM orders
+GROUP BY category
+ORDER BY avg_ship_days;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_006",
+        difficulty="hard",
+        question="Find customers whose total order value in Technology exceeds their total in all other categories combined. Show customer name, tech_total, and other_total.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+WITH category_totals AS (
+    SELECT c.customer_id, c.name,
+           SUM(CASE WHEN o.category='Technology' THEN o.amount ELSE 0 END) as tech_total,
+           SUM(CASE WHEN o.category!='Technology' THEN o.amount ELSE 0 END) as other_total
+    FROM customers c
+    JOIN orders o ON c.customer_id = o.customer_id
+    GROUP BY c.customer_id, c.name
+)
+SELECT name, tech_total, other_total
+FROM category_totals
+WHERE tech_total > other_total
+ORDER BY tech_total DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_007",
+        difficulty="hard",
+        question="Using LAG window function, for each customer show each order amount and the difference from their previous order amount. Show name, order_date, amount, prev_amount, and amount_change.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT c.name, o.order_date, o.amount,
+       LAG(o.amount) OVER (PARTITION BY c.customer_id ORDER BY o.order_date) as prev_amount,
+       o.amount - LAG(o.amount) OVER (PARTITION BY c.customer_id ORDER BY o.order_date) as amount_change
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+ORDER BY c.name, o.order_date;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_008",
+        difficulty="hard",
+        question="Identify the top 2 customers by total order amount in each customer segment. Show segment, customer name, total_amount, and rank within segment.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+WITH ranked AS (
+    SELECT c.segment, c.name,
+           SUM(o.amount) as total_amount,
+           RANK() OVER (PARTITION BY c.segment ORDER BY SUM(o.amount) DESC) as rnk
+    FROM customers c
+    JOIN orders o ON c.customer_id = o.customer_id
+    GROUP BY c.customer_id, c.segment, c.name
+)
+SELECT segment, name, total_amount, rnk
+FROM ranked
+WHERE rnk <= 2
+ORDER BY segment, rnk;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_009",
+        difficulty="hard",
+        question="Find the percentage of total revenue contributed by each customer. Show name and revenue_pct rounded to 2 decimals, ordered by pct descending.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT c.name,
+       ROUND(SUM(o.amount) * 100.0 / (SELECT SUM(amount) FROM orders), 2) as revenue_pct
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+GROUP BY c.customer_id, c.name
+ORDER BY revenue_pct DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_010",
+        difficulty="hard",
+        question="For each category, show total revenue, average order size, and count of unique customers who ordered in that category.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT o.category,
+       ROUND(SUM(o.amount), 2) as total_revenue,
+       ROUND(AVG(o.amount), 2) as avg_order_size,
+       COUNT(DISTINCT o.customer_id) as unique_customers
+FROM orders o
+GROUP BY o.category
+ORDER BY total_revenue DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_011",
+        difficulty="hard",
+        question="Find customers who have ordered in ALL three categories (Technology, Furniture, Office Supplies). Show customer name.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT c.name
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+WHERE o.category IN ('Technology','Furniture','Office Supplies')
+GROUP BY c.customer_id, c.name
+HAVING COUNT(DISTINCT o.category) = 3
+ORDER BY c.name;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_012",
+        difficulty="hard",
+        question="For each order, show the order amount as a percentile rank (0 to 1) among all orders. Show order_id, amount, and percentile_rank rounded to 2 decimals.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT order_id, amount,
+       ROUND(PERCENT_RANK() OVER (ORDER BY amount), 2) as percentile_rank
+FROM orders
+ORDER BY amount;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_013",
+        difficulty="hard",
+        question="Find orders where the shipping time was above the average shipping time for their category. Show order_id, category, ship_days, and category_avg_days (both rounded to 1 decimal).",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+WITH ship_times AS (
+    SELECT order_id, category,
+           ROUND(julianday(ship_date) - julianday(order_date), 1) as ship_days,
+           ROUND(AVG(julianday(ship_date) - julianday(order_date)) OVER (PARTITION BY category), 1) as category_avg_days
+    FROM orders
+)
+SELECT order_id, category, ship_days, category_avg_days
+FROM ship_times
+WHERE ship_days > category_avg_days
+ORDER BY ship_days DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_014",
+        difficulty="hard",
+        question="Find the first and last order date for each customer, and number of days between them (customer tenure in days). Show name, first_order, last_order, tenure_days.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+SELECT c.name,
+       MIN(o.order_date) as first_order,
+       MAX(o.order_date) as last_order,
+       CAST(julianday(MAX(o.order_date)) - julianday(MIN(o.order_date)) AS INTEGER) as tenure_days
+FROM customers c
+JOIN orders o ON c.customer_id = o.customer_id
+GROUP BY c.customer_id, c.name
+ORDER BY tenure_days DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+
+    SQLTask(
+        task_id="hard_015",
+        difficulty="hard",
+        question="Using a CTE, find customers with above-average order count AND above-average total spend. Show name, order_count, and total_spend.",
+        schema_sql=HARD_SCHEMA,
+        seed_sql=HARD_SEED,
+        solution_sql="""
+WITH customer_stats AS (
+    SELECT c.customer_id, c.name,
+           COUNT(o.order_id) as order_count,
+           SUM(o.amount) as total_spend
+    FROM customers c
+    JOIN orders o ON c.customer_id = o.customer_id
+    GROUP BY c.customer_id, c.name
+),
+averages AS (
+    SELECT AVG(order_count) as avg_orders, AVG(total_spend) as avg_spend
+    FROM customer_stats
+)
+SELECT cs.name, cs.order_count, cs.total_spend
+FROM customer_stats cs, averages
+WHERE cs.order_count > averages.avg_orders
+  AND cs.total_spend > averages.avg_spend
+ORDER BY cs.total_spend DESC;""",
+        schema_description="customers(customer_id, name, city, segment), orders(order_id, customer_id, order_date, ship_date, category, amount), returns(return_id, order_id, return_date, reason)",
+    ),
+]
+
+
+# ─────────────────────────────────────────────
+# EXPERT TASKS — Correlated subqueries, complex analytics
+# ─────────────────────────────────────────────
+
+EXPERT_SCHEMA = """
+CREATE TABLE accounts (
+    account_id INTEGER PRIMARY KEY,
+    holder_name TEXT,
+    account_type TEXT,
+    city TEXT,
+    opened_date TEXT,
+    balance REAL
+);
+CREATE TABLE transactions (
+    txn_id INTEGER PRIMARY KEY,
+    account_id INTEGER,
+    txn_date TEXT,
+    txn_type TEXT,
+    amount REAL,
+    category TEXT,
+    FOREIGN KEY (account_id) REFERENCES accounts(account_id)
+);
+CREATE TABLE loans (
+    loan_id INTEGER PRIMARY KEY,
+    account_id INTEGER,
+    loan_type TEXT,
+    principal REAL,
+    interest_rate REAL,
+    start_date TEXT,
+    status TEXT,
+    FOREIGN KEY (account_id) REFERENCES accounts(account_id)
+);"""
+
+EXPERT_SEED = """
+INSERT INTO accounts VALUES
+(1,'Arjun Sharma','Savings','Mumbai','2018-03-15',125000.00),
+(2,'Priya Patel','Current','Delhi','2016-07-22',450000.00),
+(3,'Rahul Gupta','Savings','Bangalore','2020-01-10',78000.00),
+(4,'Sneha Reddy','Current','Hyderabad','2015-11-30',890000.00),
+(5,'Vikram Singh','Savings','Chennai','2019-05-14',42000.00),
+(6,'Ananya Joshi','Savings','Pune','2021-08-09',195000.00),
+(7,'Karan Mehta','Current','Mumbai','2017-02-28',320000.00),
+(8,'Divya Nair','Savings','Bangalore','2022-04-17',55000.00);
+
+INSERT INTO transactions VALUES
+(1,1,'2024-01-05','credit',50000,'salary'),
+(2,1,'2024-01-12','debit',15000,'rent'),
+(3,1,'2024-01-20','debit',8000,'groceries'),
+(4,1,'2024-02-05','credit',50000,'salary'),
+(5,1,'2024-02-18','debit',22000,'utilities'),
+(6,2,'2024-01-08','credit',200000,'business'),
+(7,2,'2024-01-15','debit',80000,'supplier'),
+(8,2,'2024-01-28','debit',45000,'rent'),
+(9,2,'2024-02-10','credit',175000,'business'),
+(10,2,'2024-02-22','debit',60000,'supplier'),
+(11,3,'2024-01-10','credit',45000,'salary'),
+(12,3,'2024-01-25','debit',12000,'rent'),
+(13,3,'2024-02-10','credit',45000,'salary'),
+(14,3,'2024-02-20','debit',18000,'utilities'),
+(15,4,'2024-01-03','credit',500000,'business'),
+(16,4,'2024-01-20','debit',200000,'investment'),
+(17,4,'2024-02-05','credit',480000,'business'),
+(18,4,'2024-02-25','debit',150000,'supplier'),
+(19,5,'2024-01-15','credit',38000,'salary'),
+(20,5,'2024-01-30','debit',15000,'rent'),
+(21,6,'2024-01-08','credit',85000,'salary'),
+(22,6,'2024-01-22','debit',25000,'rent'),
+(23,6,'2024-02-08','credit',85000,'salary'),
+(24,6,'2024-02-20','debit',30000,'investment'),
+(25,7,'2024-01-12','credit',150000,'business'),
+(26,7,'2024-01-28','debit',75000,'supplier'),
+(27,7,'2024-02-15','credit',160000,'business'),
+(28,8,'2024-01-18','credit',52000,'salary'),
+(29,8,'2024-02-18','credit',52000,'salary'),
+(30,8,'2024-02-25','debit',20000,'rent');
+
+INSERT INTO loans VALUES
+(1,1,'Home Loan',2500000,8.5,'2020-06-15','active'),
+(2,2,'Business Loan',5000000,10.2,'2019-03-20','active'),
+(3,3,'Personal Loan',150000,14.5,'2022-08-10','active'),
+(4,5,'Vehicle Loan',800000,9.8,'2021-12-05','active'),
+(5,6,'Home Loan',3500000,8.2,'2022-09-15','active'),
+(6,7,'Business Loan',2000000,11.5,'2018-04-22','closed'),
+(7,8,'Personal Loan',100000,15.2,'2023-01-30','active');"""
+
+EXPERT_TASKS = [
+    SQLTask(
+        task_id="expert_001",
+        difficulty="expert",
+        question="For each account, calculate the net cash flow (total credits minus total debits) per month in 2024. Show holder_name, month (YYYY-MM), and net_flow. Order by holder_name, month.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+SELECT a.holder_name,
+       SUBSTR(t.txn_date,1,7) as month,
+       ROUND(SUM(CASE WHEN t.txn_type='credit' THEN t.amount ELSE -t.amount END), 2) as net_flow
+FROM accounts a
+JOIN transactions t ON a.account_id = t.account_id
+WHERE t.txn_date LIKE '2024%'
+GROUP BY a.account_id, a.holder_name, month
+ORDER BY a.holder_name, month;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_002",
+        difficulty="expert",
+        question="Find accounts where the total loan principal exceeds 5x their current account balance. Show holder_name, balance, total_loan_principal, and the ratio (rounded to 2 decimals).",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+SELECT a.holder_name, a.balance,
+       SUM(l.principal) as total_loan_principal,
+       ROUND(SUM(l.principal) / a.balance, 2) as debt_to_balance_ratio
+FROM accounts a
+JOIN loans l ON a.account_id = l.account_id
+WHERE l.status = 'active'
+GROUP BY a.account_id, a.holder_name, a.balance
+HAVING SUM(l.principal) > 5 * a.balance
+ORDER BY debt_to_balance_ratio DESC;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_003",
+        difficulty="expert",
+        question="Using a CTE chain, find the top spending category per account (by total debit amount). Show holder_name, top_category, and category_spend.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+WITH category_spend AS (
+    SELECT a.account_id, a.holder_name, t.category,
+           SUM(t.amount) as spend,
+           RANK() OVER (PARTITION BY a.account_id ORDER BY SUM(t.amount) DESC) as rnk
+    FROM accounts a
+    JOIN transactions t ON a.account_id = t.account_id
+    WHERE t.txn_type = 'debit'
+    GROUP BY a.account_id, a.holder_name, t.category
+)
+SELECT holder_name, category as top_category, spend as category_spend
+FROM category_spend
+WHERE rnk = 1
+ORDER BY category_spend DESC;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_004",
+        difficulty="expert",
+        question="Calculate the estimated annual interest cost for each active loan. Show holder_name, loan_type, principal, interest_rate, and annual_interest_cost (principal * interest_rate / 100), ordered by annual_interest_cost descending.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+SELECT a.holder_name, l.loan_type, l.principal, l.interest_rate,
+       ROUND(l.principal * l.interest_rate / 100, 2) as annual_interest_cost
+FROM accounts a
+JOIN loans l ON a.account_id = l.account_id
+WHERE l.status = 'active'
+ORDER BY annual_interest_cost DESC;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_005",
+        difficulty="expert",
+        question="For each city, calculate the average account balance, total transaction volume, and number of active loans. Show city, avg_balance, total_txn_volume, and active_loan_count.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+SELECT a.city,
+       ROUND(AVG(a.balance), 2) as avg_balance,
+       ROUND(SUM(t.amount), 2) as total_txn_volume,
+       COUNT(DISTINCT l.loan_id) as active_loan_count
+FROM accounts a
+LEFT JOIN transactions t ON a.account_id = t.account_id
+LEFT JOIN loans l ON a.account_id = l.account_id AND l.status = 'active'
+GROUP BY a.city
+ORDER BY avg_balance DESC;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_006",
+        difficulty="expert",
+        question="Find accounts that received salary credits in both January and February 2024. Show holder_name.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+SELECT a.holder_name
+FROM accounts a
+WHERE EXISTS (
+    SELECT 1 FROM transactions t
+    WHERE t.account_id = a.account_id
+      AND t.txn_type = 'credit'
+      AND t.category = 'salary'
+      AND t.txn_date LIKE '2024-01%'
+)
+AND EXISTS (
+    SELECT 1 FROM transactions t
+    WHERE t.account_id = a.account_id
+      AND t.txn_type = 'credit'
+      AND t.category = 'salary'
+      AND t.txn_date LIKE '2024-02%'
+)
+ORDER BY a.holder_name;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_007",
+        difficulty="expert",
+        question="Calculate savings rate for each salary earner: (total_credits - total_debits) / total_credits. Show holder_name and savings_rate as percentage rounded to 1 decimal. Only include accounts with salary transactions.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+WITH flows AS (
+    SELECT a.account_id, a.holder_name,
+           SUM(CASE WHEN t.txn_type='credit' THEN t.amount ELSE 0 END) as total_credits,
+           SUM(CASE WHEN t.txn_type='debit' THEN t.amount ELSE 0 END) as total_debits
+    FROM accounts a
+    JOIN transactions t ON a.account_id = t.account_id
+    GROUP BY a.account_id, a.holder_name
+)
+SELECT f.holder_name,
+       ROUND((f.total_credits - f.total_debits) * 100.0 / f.total_credits, 1) as savings_rate
+FROM flows f
+WHERE f.account_id IN (
+    SELECT DISTINCT account_id FROM transactions WHERE category = 'salary'
+)
+ORDER BY savings_rate DESC;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_008",
+        difficulty="expert",
+        question="Find the account with the highest transaction frequency (most transactions per month on average). Show holder_name, total_transactions, months_active (distinct months with transactions), and avg_txn_per_month rounded to 2 decimals.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+SELECT a.holder_name,
+       COUNT(t.txn_id) as total_transactions,
+       COUNT(DISTINCT SUBSTR(t.txn_date,1,7)) as months_active,
+       ROUND(CAST(COUNT(t.txn_id) AS REAL) / COUNT(DISTINCT SUBSTR(t.txn_date,1,7)), 2) as avg_txn_per_month
+FROM accounts a
+JOIN transactions t ON a.account_id = t.account_id
+GROUP BY a.account_id, a.holder_name
+ORDER BY avg_txn_per_month DESC
+LIMIT 1;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_009",
+        difficulty="expert",
+        question="Identify accounts that have both an active loan AND whose total debits in 2024 exceed their total credits in 2024 (cash flow negative). Show holder_name, net_flow_2024, and total_active_loan_principal.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+WITH cash_flow AS (
+    SELECT account_id,
+           SUM(CASE WHEN txn_type='credit' THEN amount ELSE -amount END) as net_flow
+    FROM transactions
+    WHERE txn_date LIKE '2024%'
+    GROUP BY account_id
+),
+active_loans AS (
+    SELECT account_id, SUM(principal) as total_principal
+    FROM loans WHERE status='active'
+    GROUP BY account_id
+)
+SELECT a.holder_name, ROUND(cf.net_flow, 2) as net_flow_2024, al.total_principal
+FROM accounts a
+JOIN cash_flow cf ON a.account_id = cf.account_id
+JOIN active_loans al ON a.account_id = al.account_id
+WHERE cf.net_flow < 0
+ORDER BY cf.net_flow;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+
+    SQLTask(
+        task_id="expert_010",
+        difficulty="expert",
+        question="Create a financial health score for each account: score = (balance/100000)*40 + (net_credits_2024/50000)*40 - (active_loan_principal/1000000)*20, capped between 0 and 100. Show holder_name and health_score rounded to 1 decimal, ordered by score descending.",
+        schema_sql=EXPERT_SCHEMA,
+        seed_sql=EXPERT_SEED,
+        solution_sql="""
+WITH metrics AS (
+    SELECT a.account_id, a.holder_name, a.balance,
+           COALESCE(SUM(CASE WHEN t.txn_type='credit' THEN t.amount ELSE 0 END), 0) as net_credits,
+           COALESCE((SELECT SUM(principal) FROM loans l WHERE l.account_id=a.account_id AND l.status='active'), 0) as loan_principal
+    FROM accounts a
+    LEFT JOIN transactions t ON a.account_id = t.account_id AND t.txn_date LIKE '2024%'
+    GROUP BY a.account_id, a.holder_name, a.balance
+)
+SELECT holder_name,
+       ROUND(MIN(100, MAX(0,
+           (balance/100000.0)*40 +
+           (net_credits/50000.0)*40 -
+           (loan_principal/1000000.0)*20
+       )), 1) as health_score
+FROM metrics
+ORDER BY health_score DESC;""",
+        schema_description="accounts(account_id, holder_name, account_type, city, opened_date, balance), transactions(txn_id, account_id, txn_date, txn_type[credit/debit], amount, category), loans(loan_id, account_id, loan_type, principal, interest_rate, start_date, status)",
+    ),
+]
+
+
+# ─────────────────────────────────────────────
+# Task registry
+# ─────────────────────────────────────────────
+
+ALL_TASKS: Dict[str, SQLTask] = {}
+for _t in EASY_TASKS + MEDIUM_TASKS + HARD_TASKS + EXPERT_TASKS:
+    ALL_TASKS[_t.task_id] = _t
+
+TASKS_BY_DIFFICULTY: Dict[str, List[SQLTask]] = {
+    "easy": EASY_TASKS,
+    "medium": MEDIUM_TASKS,
+    "hard": HARD_TASKS,
+    "expert": EXPERT_TASKS,
+}
+
+def get_task(task_id: str) -> Optional[SQLTask]:
+    return ALL_TASKS.get(task_id)
+
+def get_tasks_by_difficulty(difficulty: str) -> List[SQLTask]:
+    return TASKS_BY_DIFFICULTY.get(difficulty, [])