SQL Query OpenEnv Environment

.dockerignore

@@ -0,0 +1,7 @@
+__pycache__
+*.pyc
+.git
+.gitignore
+outputs/
+*.egg-info
+.venv

Dockerfile

@@ -0,0 +1,23 @@
+FROM python:3.11-slim
+
+# Create non-root user (HF Spaces runs with user ID 1000)
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+WORKDIR /app
+
+# Install dependencies
+COPY --chown=user server/requirements.txt .
+RUN pip install --no-cache-dir --user -r requirements.txt
+
+# Copy all source code
+COPY --chown=user . .
+
+# HF Spaces uses port 7860 by default
+ENV PORT=7860
+EXPOSE 7860
+
+# Run the FastAPI server
+CMD ["python", "-m", "uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,11 +1,35 @@
 ---
-title: Sql Query Env
-emoji: 🏢
-colorFrom: yellow
-colorTo: red
+title: SQL Query Environment
+emoji: 🗃️
+colorFrom: blue
+colorTo: green
 sdk: docker
+app_port: 7860
+tags:
+  - openenv
+  - sql
+  - text-to-sql
+  - reinforcement-learning
 pinned: false
-license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SQL Query Environment
+
+An OpenEnv-compatible environment where AI agents learn to translate natural language questions into correct SQL queries.
+
+## API Endpoints
+
+- `GET /health` — Health check
+- `POST /reset` — Start new episode: `{"task_id": "task_1"}`
+- `POST /step` — Submit SQL: `{"task_id": "task_1", "sql_query": "SELECT ..."}`
+- `GET /state` — Current episode state
+
+## Tasks
+
+| ID | Difficulty | Description |
+|----|-----------|-------------|
+| task_1 | Easy | Single-table filtering + ordering |
+| task_2 | Medium | JOIN + GROUP BY + HAVING |
+| task_3 | Hard | Subquery + multi-JOIN + aggregation |
+
+Grading is fully deterministic with partial credit (0.0 to 1.0).

__init__.py

@@ -0,0 +1,3 @@
+from .models import SQLAction, SQLObservation
+
+__all__ = ["SQLAction", "SQLObservation"]

inference.py

@@ -0,0 +1,193 @@
+"""
+Baseline inference script for the SQL Query Environment.
+
+Uses the OpenAI-compatible API to run an LLM agent against all 3 tasks.
+Reads credentials from environment variables:
+  - API_BASE_URL: The API endpoint for the LLM
+  - MODEL_NAME:   The model identifier to use
+  - HF_TOKEN:     Your Hugging Face / API key
+
+Usage:
+    API_BASE_URL=https://router.huggingface.co/v1 \
+    MODEL_NAME=Qwen/Qwen2.5-72B-Instruct \
+    HF_TOKEN=hf_xxx \
+    python inference.py
+"""
+
+import os
+import sys
+import json
+import time
+import requests
+
+from openai import OpenAI
+
+# ── Configuration ──
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY", "")
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+
+# Environment URL (local or HF Space)
+ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")
+
+MAX_ATTEMPTS = 3  # Max steps per task
+TASKS = ["task_1", "task_2", "task_3"]
+
+
+def call_environment(endpoint: str, payload: dict = None) -> dict:
+    """Call the environment's HTTP API."""
+    url = f"{ENV_URL.rstrip('/')}/{endpoint.lstrip('/')}"
+    if payload is not None:
+        resp = requests.post(url, json=payload, timeout=30)
+    else:
+        resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    return resp.json()
+
+
+def generate_sql(client: OpenAI, schema: str, question: str, history: str = "") -> str:
+    """Ask the LLM to generate a SQL query."""
+    system_prompt = (
+        "You are an expert SQL query writer. Given a database schema and a "
+        "natural language question, write a single SQLite-compatible SQL query "
+        "that answers the question.\n\n"
+        "RULES:\n"
+        "- Return ONLY the SQL query, nothing else.\n"
+        "- Do NOT include markdown code fences, explanations, or comments.\n"
+        "- Use proper SQLite syntax.\n"
+        "- Pay attention to column names and table relationships.\n"
+        "- When asked to round values, use ROUND(value, decimal_places).\n"
+        "- Use single quotes for string literals.\n"
+    )
+
+    user_prompt = f"DATABASE SCHEMA:\n{schema}\n\nQUESTION: {question}"
+    if history:
+        user_prompt += f"\n\nPREVIOUS ATTEMPTS AND FEEDBACK:\n{history}"
+        user_prompt += "\n\nPlease fix your query based on the feedback above."
+
+    try:
+        completion = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt},
+            ],
+            temperature=0.1,
+            max_tokens=500,
+        )
+        response = completion.choices[0].message.content or ""
+        # Clean up: remove markdown fences if present
+        response = response.strip()
+        if response.startswith("```sql"):
+            response = response[6:]
+        if response.startswith("```"):
+            response = response[3:]
+        if response.endswith("```"):
+            response = response[:-3]
+        return response.strip()
+    except Exception as e:
+        print(f"  LLM API error: {e}")
+        return "SELECT 1"
+
+
+def run_task(client: OpenAI, task_id: str) -> float:
+    """Run the agent on a single task. Returns the best reward achieved."""
+    print(f"\n{'='*60}")
+    print(f"TASK: {task_id.upper()}")
+    print(f"{'='*60}")
+
+    # Reset environment for this task
+    reset_resp = call_environment("/reset", {"task_id": task_id})
+    obs = reset_resp["observation"]
+
+    print(f"Question: {obs['task_description']}")
+    print(f"Difficulty: {obs['difficulty']}")
+
+    best_reward = 0.0
+    history = ""
+
+    for attempt in range(1, MAX_ATTEMPTS + 1):
+        print(f"\n--- Attempt {attempt}/{MAX_ATTEMPTS} ---")
+
+        # Generate SQL query
+        sql = generate_sql(
+            client,
+            schema=obs["schema_description"],
+            question=obs["task_description"],
+            history=history,
+        )
+        print(f"SQL: {sql[:200]}{'...' if len(sql) > 200 else ''}")
+
+        # Submit to environment
+        step_resp = call_environment("/step", {
+            "task_id": task_id,
+            "sql_query": sql,
+        })
+        obs = step_resp["observation"]
+        reward = step_resp["reward"]
+        done = step_resp["done"]
+
+        print(f"Reward: {reward:.2f}")
+        print(f"Feedback: {obs['feedback']}")
+
+        best_reward = max(best_reward, reward)
+
+        # Track history for retry
+        history += f"\nAttempt {attempt}: SQL: {sql}\n"
+        history += f"  Reward: {reward}, Feedback: {obs['feedback']}\n"
+        if obs.get("query_error"):
+            history += f"  Error: {obs['query_error']}\n"
+
+        if done:
+            break
+
+    print(f"\nBest reward for {task_id}: {best_reward:.2f}")
+    return best_reward
+
+
+def main():
+    """Run the baseline agent on all tasks and report scores."""
+    print("=" * 60)
+    print("SQL Query Environment - Baseline Inference")
+    print("=" * 60)
+    print(f"API URL:  {API_BASE_URL}")
+    print(f"Model:    {MODEL_NAME}")
+    print(f"Env URL:  {ENV_URL}")
+    print()
+
+    # Verify environment is up
+    try:
+        health = call_environment("/health")
+        print(f"Environment health: {health}")
+    except Exception as e:
+        print(f"ERROR: Cannot reach environment at {ENV_URL}: {e}")
+        print("Make sure the environment server is running.")
+        sys.exit(1)
+
+    # Create OpenAI client
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+
+    # Run all tasks
+    results = {}
+    total_score = 0.0
+
+    for task_id in TASKS:
+        score = run_task(client, task_id)
+        results[task_id] = score
+        total_score += score
+
+    # Summary
+    print("\n" + "=" * 60)
+    print("FINAL RESULTS")
+    print("=" * 60)
+    for task_id, score in results.items():
+        print(f"  {task_id:8s}: {score:.2f}")
+    avg_score = total_score / len(TASKS)
+    print(f"  {'AVERAGE':8s}: {avg_score:.2f}")
+    print("=" * 60)
+
+    return results
+
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,61 @@
+"""
+Typed Action and Observation models for the SQL Query Environment.
+
+Action: The agent submits a SQL query string and a task_id.
+Observation: The environment returns schema info, query results, feedback, and reward.
+"""
+
+from pydantic import Field
+from openenv.core.env_server.types import Action, Observation
+
+
+class SQLAction(Action):
+    """Action submitted by the agent: a SQL query to execute."""
+
+    task_id: str = Field(
+        ...,
+        description="ID of the task being attempted (task_1, task_2, task_3)",
+    )
+    sql_query: str = Field(
+        ...,
+        description="The SQL query string to execute against the database",
+    )
+
+
+class SQLObservation(Observation):
+    """Observation returned to the agent after each step."""
+
+    # Task information
+    task_id: str = Field(default="", description="Current task ID")
+    task_description: str = Field(
+        default="", description="Natural language question the agent must answer"
+    )
+    difficulty: str = Field(default="", description="easy, medium, or hard")
+
+    # Database schema
+    schema_description: str = Field(
+        default="", description="SQL CREATE TABLE statements describing the database"
+    )
+
+    # Query result / feedback
+    query_result: str = Field(
+        default="",
+        description="Result of the executed SQL query (rows as text), or error message",
+    )
+    query_error: bool = Field(
+        default=False, description="True if the SQL query caused an error"
+    )
+    feedback: str = Field(
+        default="",
+        description="Human-readable feedback on the query result",
+    )
+
+    # Scoring
+    reward: float = Field(default=0.0, description="Score from 0.0 to 1.0")
+    done: bool = Field(default=False, description="True if the episode is complete")
+
+    # Metadata
+    step_count: int = Field(default=0, description="Number of steps taken so far")
+    max_steps: int = Field(
+        default=3, description="Maximum steps allowed per task"
+    )

openenv.yaml

@@ -0,0 +1,13 @@
+name: sql_query_env
+version: "1.0.0"
+description: >
+  A real-world SQL query environment where AI agents learn to translate
+  natural language questions into correct SQL queries. Features 3 tasks
+  with increasing difficulty (easy/medium/hard), deterministic grading
+  via SQLite execution, and partial-credit reward shaping.
+tags:
+  - openenv
+  - sql
+  - text-to-sql
+  - data-analysis
+  - natural-language-processing

pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.backends._legacy:_Backend"
+
+[project]
+name = "sql_query_env"
+version = "1.0.0"
+description = "SQL Query OpenEnv Environment - Train agents to write SQL from natural language"
+requires-python = ">=3.10"
+dependencies = [
+    "openenv-core>=0.2.1",
+    "fastapi>=0.104.0",
+    "uvicorn>=0.24.0",
+    "pydantic>=2.0.0",
+    "websockets>=12.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest", "httpx"]
+
+[tool.setuptools.packages.find]
+include = ["sql_query_env*"]

server/__init__.py

@@ -0,0 +1 @@
+"""SQL Query Environment Server."""

server/app.py

@@ -0,0 +1,104 @@
+"""FastAPI server for the SQL Query Environment.
+
+Exposes HTTP endpoints:
+  POST /reset   - Start a new episode
+  POST /step    - Submit a SQL query
+  GET  /state   - Get current state
+  GET  /health  - Health check
+"""
+
+import os
+import sys
+from typing import Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+# Add parent + server to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from models import SQLAction, SQLObservation
+from sql_environment import SQLQueryEnvironment
+
+
+# ── Request/Response models ──
+
+class ResetRequest(BaseModel):
+    task_id: Optional[str] = "task_1"
+
+
+class StepRequest(BaseModel):
+    task_id: str = "task_1"
+    sql_query: str = ""
+
+
+# ── Environment instance ──
+env = SQLQueryEnvironment()
+
+
+app = FastAPI(
+    title="SQL Query Environment",
+    description=(
+        "An OpenEnv environment where AI agents learn to write SQL queries. "
+        "Features 3 tasks with increasing difficulty and deterministic grading."
+    ),
+    version="1.0.0",
+)
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {"status": "healthy"}
+
+
+@app.post("/reset")
+async def reset(request: ResetRequest = ResetRequest()):
+    """Reset the environment and start a new episode."""
+    task_id = request.task_id or "task_1"
+    obs = env.reset(task_id=task_id)
+    return {
+        "observation": obs.model_dump(),
+        "reward": 0.0,
+        "done": False,
+        "info": {"episode_started": True, "task_id": task_id},
+    }
+
+
+@app.post("/step")
+async def step(request: StepRequest):
+    """Submit a SQL query and get the result with grading."""
+    if not request.sql_query:
+        raise HTTPException(status_code=400, detail="sql_query is required")
+
+    action = SQLAction(task_id=request.task_id, sql_query=request.sql_query)
+    obs = env.step(action)
+
+    return {
+        "observation": obs.model_dump(),
+        "reward": obs.reward,
+        "done": obs.done,
+        "info": {
+            "step_count": obs.step_count,
+            "feedback": obs.feedback,
+        },
+    }
+
+
+@app.get("/state")
+async def get_state():
+    """Get current episode state."""
+    s = env.state
+    return {
+        "episode_id": s.episode_id,
+        "step_count": s.step_count,
+    }
+
+
+# ── Run with uvicorn ──
+if __name__ == "__main__":
+    import uvicorn
+
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)

server/database.py

@@ -0,0 +1,144 @@
+"""
+Database setup for the SQL Query Environment.
+
+Creates an in-memory SQLite database with three tables:
+  - departments: id, name, budget, location
+  - employees: id, name, department_id, salary, hire_date, is_active
+  - projects: id, name, department_id, lead_employee_id, budget, status, start_date
+
+All data is deterministic so grading is reproducible.
+"""
+
+import sqlite3
+from typing import Optional
+
+SCHEMA_SQL = """
+CREATE TABLE departments (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    budget REAL NOT NULL,
+    location TEXT NOT NULL
+);
+
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department_id INTEGER NOT NULL,
+    salary REAL NOT NULL,
+    hire_date TEXT NOT NULL,
+    is_active INTEGER NOT NULL DEFAULT 1,
+    FOREIGN KEY (department_id) REFERENCES departments(id)
+);
+
+CREATE TABLE projects (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department_id INTEGER NOT NULL,
+    lead_employee_id INTEGER NOT NULL,
+    budget REAL NOT NULL,
+    status TEXT NOT NULL CHECK(status IN ('active', 'completed', 'cancelled')),
+    start_date TEXT NOT NULL,
+    FOREIGN KEY (department_id) REFERENCES departments(id),
+    FOREIGN KEY (lead_employee_id) REFERENCES employees(id)
+);
+"""
+
+SEED_SQL = """
+INSERT INTO departments VALUES (1, 'Engineering', 500000.00, 'Bangalore');
+INSERT INTO departments VALUES (2, 'Marketing', 200000.00, 'Mumbai');
+INSERT INTO departments VALUES (3, 'Sales', 300000.00, 'Delhi');
+INSERT INTO departments VALUES (4, 'HR', 150000.00, 'Bangalore');
+INSERT INTO departments VALUES (5, 'Finance', 250000.00, 'Mumbai');
+
+INSERT INTO employees VALUES (1,  'Arjun Sharma',    1, 85000.00, '2020-03-15', 1);
+INSERT INTO employees VALUES (2,  'Priya Patel',     1, 92000.00, '2019-07-01', 1);
+INSERT INTO employees VALUES (3,  'Rahul Verma',     1, 78000.00, '2021-01-10', 1);
+INSERT INTO employees VALUES (4,  'Sneha Gupta',     2, 65000.00, '2020-06-20', 1);
+INSERT INTO employees VALUES (5,  'Vikram Singh',    2, 70000.00, '2018-11-05', 1);
+INSERT INTO employees VALUES (6,  'Anita Desai',     3, 72000.00, '2019-09-12', 1);
+INSERT INTO employees VALUES (7,  'Karan Mehta',     3, 68000.00, '2021-04-01', 1);
+INSERT INTO employees VALUES (8,  'Deepa Nair',      3, 75000.00, '2020-02-28', 0);
+INSERT INTO employees VALUES (9,  'Suresh Kumar',    4, 60000.00, '2022-01-15', 1);
+INSERT INTO employees VALUES (10, 'Meera Joshi',     4, 58000.00, '2021-08-20', 1);
+INSERT INTO employees VALUES (11, 'Amit Rao',        5, 88000.00, '2019-05-10', 1);
+INSERT INTO employees VALUES (12, 'Lakshmi Iyer',    5, 82000.00, '2020-10-01', 1);
+INSERT INTO employees VALUES (13, 'Ravi Krishnan',   1, 95000.00, '2018-03-20', 1);
+INSERT INTO employees VALUES (14, 'Pooja Reddy',     2, 62000.00, '2022-06-15', 1);
+INSERT INTO employees VALUES (15, 'Nikhil Agarwal',  3, 71000.00, '2020-12-01', 1);
+
+INSERT INTO projects VALUES (1, 'Cloud Migration',     1, 2,  120000.00, 'active',    '2024-01-15');
+INSERT INTO projects VALUES (2, 'Mobile App v2',       1, 1,  80000.00,  'active',    '2024-03-01');
+INSERT INTO projects VALUES (3, 'Brand Refresh',       2, 5,  45000.00,  'completed', '2023-06-01');
+INSERT INTO projects VALUES (4, 'Q4 Campaign',         2, 4,  60000.00,  'active',    '2024-09-01');
+INSERT INTO projects VALUES (5, 'CRM Integration',     3, 6,  90000.00,  'active',    '2024-02-15');
+INSERT INTO projects VALUES (6, 'Sales Dashboard',     3, 15, 35000.00,  'completed', '2023-11-01');
+INSERT INTO projects VALUES (7, 'Payroll Automation',   4, 9,  50000.00,  'active',    '2024-04-01');
+INSERT INTO projects VALUES (8, 'Annual Audit Tool',    5, 11, 70000.00,  'cancelled', '2024-01-10');
+INSERT INTO projects VALUES (9, 'Data Pipeline',        1, 13, 150000.00, 'active',    '2024-06-01');
+INSERT INTO projects VALUES (10, 'Employee Portal',     4, 10, 40000.00,  'completed', '2023-09-15');
+"""
+
+SCHEMA_DESCRIPTION = """Tables in the database:
+
+CREATE TABLE departments (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    budget REAL NOT NULL,
+    location TEXT NOT NULL
+);
+
+CREATE TABLE employees (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department_id INTEGER NOT NULL,
+    salary REAL NOT NULL,
+    hire_date TEXT NOT NULL,       -- format: YYYY-MM-DD
+    is_active INTEGER NOT NULL,    -- 1 = active, 0 = inactive
+    FOREIGN KEY (department_id) REFERENCES departments(id)
+);
+
+CREATE TABLE projects (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    department_id INTEGER NOT NULL,
+    lead_employee_id INTEGER NOT NULL,
+    budget REAL NOT NULL,
+    status TEXT NOT NULL,           -- 'active', 'completed', or 'cancelled'
+    start_date TEXT NOT NULL,       -- format: YYYY-MM-DD
+    FOREIGN KEY (department_id) REFERENCES departments(id),
+    FOREIGN KEY (lead_employee_id) REFERENCES employees(id)
+);""".strip()
+
+
+def create_database() -> sqlite3.Connection:
+    """Create a fresh in-memory SQLite database with schema and seed data."""
+    conn = sqlite3.connect(":memory:")
+    conn.execute("PRAGMA foreign_keys = ON;")
+    conn.executescript(SCHEMA_SQL)
+    conn.executescript(SEED_SQL)
+    conn.commit()
+    return conn
+
+
+def execute_query(
+    conn: sqlite3.Connection, sql: str
+) -> tuple[list[tuple] | None, list[str] | None, str | None]:
+    """
+    Execute a SQL query safely.
+
+    Returns:
+        (rows, column_names, error_message)
+        - On success: (rows_list, columns_list, None)
+        - On error: (None, None, error_string)
+    """
+    try:
+        cursor = conn.execute(sql)
+        if cursor.description is not None:
+            columns = [desc[0] for desc in cursor.description]
+            rows = cursor.fetchall()
+            return rows, columns, None
+        else:
+            conn.commit()
+            return None, None, None
+    except Exception as e:
+        return None, None, str(e)

server/grader.py

@@ -0,0 +1,122 @@
+"""
+Deterministic grader for the SQL Query Environment.
+
+Compares the agent's query result against the expected result.
+Produces a score between 0.0 and 1.0 with partial credit:
+
+  0.0  — SQL error or no result returned
+  0.1  — Query runs but returns wrong column count
+  0.2  — Correct column count but wrong column names
+  0.3  — Correct columns but wrong number of rows
+  0.5  — Correct structure, some rows match
+  0.7  — Most rows match (>= 70% of expected rows found)
+  0.9  — All rows match but in wrong order
+  1.0  — Exact match (columns, rows, order)
+"""
+
+from tasks import TaskDefinition
+
+
+def _normalize_value(val) -> str:
+    """Normalize a value for comparison (handle float rounding, casing)."""
+    if val is None:
+        return "NULL"
+    if isinstance(val, float):
+        return f"{val:.2f}"
+    return str(val).strip().lower()
+
+
+def _normalize_row(row: tuple) -> tuple:
+    return tuple(_normalize_value(v) for v in row)
+
+
+def grade_query(
+    task: TaskDefinition,
+    agent_rows: list[tuple] | None,
+    agent_columns: list[str] | None,
+    query_error: str | None,
+) -> tuple[float, str]:
+    """
+    Grade the agent's SQL query result against the expected result.
+
+    Returns:
+        (score, feedback_string)
+    """
+    # SQL error → 0.0
+    if query_error is not None:
+        return 0.0, f"SQL error: {query_error}"
+
+    # No result (e.g. INSERT/UPDATE instead of SELECT) → 0.0
+    if agent_rows is None or agent_columns is None:
+        return 0.0, "Query did not return any result rows. Expected a SELECT query."
+
+    expected_cols = [c.lower() for c in task.expected_columns]
+    actual_cols = [c.lower() for c in agent_columns]
+
+    # Wrong column count → 0.1
+    if len(actual_cols) != len(expected_cols):
+        return 0.1, (
+            f"Wrong number of columns. Expected {len(expected_cols)} "
+            f"({', '.join(expected_cols)}), got {len(actual_cols)} "
+            f"({', '.join(actual_cols)})."
+        )
+
+    # Wrong column names → 0.2
+    if actual_cols != expected_cols:
+        return 0.2, (
+            f"Column names don't match. Expected {expected_cols}, got {actual_cols}. "
+            f"Note: column names must match exactly."
+        )
+
+    # Columns match — now compare rows
+    expected_normalized = [_normalize_row(r) for r in task.expected_rows]
+    actual_normalized = [_normalize_row(r) for r in agent_rows]
+
+    # Wrong row count → partial credit based on matching rows
+    if len(actual_normalized) != len(expected_normalized):
+        # Count how many expected rows appear in actual
+        expected_set = set(expected_normalized)
+        actual_set = set(actual_normalized)
+        matching = len(expected_set & actual_set)
+
+        if matching == 0:
+            return 0.3, (
+                f"Correct columns but wrong number of rows. "
+                f"Expected {len(expected_normalized)} rows, got {len(actual_normalized)}. "
+                f"No matching rows found."
+            )
+        ratio = matching / len(expected_normalized)
+        score = 0.3 + ratio * 0.4  # Scale from 0.3 to 0.7
+        return round(score, 2), (
+            f"Correct columns but wrong row count. "
+            f"Expected {len(expected_normalized)} rows, got {len(actual_normalized)}. "
+            f"{matching}/{len(expected_normalized)} expected rows found."
+        )
+
+    # Same row count — check content
+    expected_set = set(expected_normalized)
+    actual_set = set(actual_normalized)
+
+    if expected_set != actual_set:
+        # Some rows match
+        matching = len(expected_set & actual_set)
+        if matching == 0:
+            return 0.5, (
+                "Correct structure (columns and row count) but row values "
+                "don't match any expected rows."
+            )
+        ratio = matching / len(expected_normalized)
+        score = 0.5 + ratio * 0.2
+        return round(score, 2), (
+            f"Partially correct. {matching}/{len(expected_normalized)} rows match."
+        )
+
+    # All rows present — check order
+    if actual_normalized != expected_normalized:
+        return 0.9, (
+            "All correct rows found but in wrong order. "
+            "Check your ORDER BY clause."
+        )
+
+    # Perfect match!
+    return 1.0, "Correct! Query returned the exact expected result."

server/requirements.txt

@@ -0,0 +1,4 @@
+fastapi>=0.104.0
+uvicorn>=0.24.0
+pydantic>=2.0.0
+openenv-core>=0.2.1

server/sql_environment.py

@@ -0,0 +1,168 @@
+"""
+SQL Query Environment — Core logic.
+
+Implements the OpenEnv Environment interface:
+  - reset()  → creates fresh DB, picks a task, returns initial observation
+  - step()   → receives SQL, executes it, grades it, returns observation
+  - state()  → returns current episode state
+"""
+
+import sqlite3
+import sys
+import os
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+sys.path.insert(0, os.path.dirname(__file__))
+from database import create_database, execute_query, SCHEMA_DESCRIPTION
+from tasks import ALL_TASKS, TASK_LIST, TaskDefinition
+from grader import grade_query
+
+# Import models
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from models import SQLAction, SQLObservation
+
+
+MAX_STEPS_PER_TASK = 3
+
+
+class SQLQueryEnvironment(Environment):
+    """
+    An environment where an AI agent writes SQL queries to answer
+    natural language questions about a company database.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._db: sqlite3.Connection | None = None
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._current_task: TaskDefinition | None = None
+        self._step_count: int = 0
+        self._done: bool = False
+        self._best_reward: float = 0.0
+
+    def reset(self, **kwargs) -> SQLObservation:
+        """Reset environment: fresh DB, pick task, return initial observation."""
+        if self._db is not None:
+            self._db.close()
+        self._db = create_database()
+
+        task_id = kwargs.get("task_id", None)
+        if task_id and task_id in ALL_TASKS:
+            self._current_task = ALL_TASKS[task_id]
+        else:
+            self._current_task = TASK_LIST[0]
+
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._step_count = 0
+        self._done = False
+        self._best_reward = 0.0
+
+        return SQLObservation(
+            task_id=self._current_task.task_id,
+            task_description=self._current_task.description,
+            difficulty=self._current_task.difficulty,
+            schema_description=SCHEMA_DESCRIPTION,
+            query_result="",
+            query_error=False,
+            feedback="Environment reset. Submit a SQL query to answer the question.",
+            reward=0.0,
+            done=False,
+            step_count=0,
+            max_steps=MAX_STEPS_PER_TASK,
+        )
+
+    def step(self, action: SQLAction) -> SQLObservation:
+        """Execute agent's SQL, grade it, return observation."""
+        if self._done:
+            return SQLObservation(
+                task_id=self._current_task.task_id if self._current_task else "",
+                task_description="",
+                difficulty="",
+                schema_description=SCHEMA_DESCRIPTION,
+                query_result="",
+                query_error=False,
+                feedback="Episode is already complete. Call reset() to start a new one.",
+                reward=0.0,
+                done=True,
+                step_count=self._step_count,
+                max_steps=MAX_STEPS_PER_TASK,
+            )
+
+        self._step_count += 1
+        self._state.step_count = self._step_count
+
+        # Validate task_id
+        if action.task_id not in ALL_TASKS:
+            return SQLObservation(
+                task_id=action.task_id,
+                task_description="",
+                difficulty="",
+                schema_description=SCHEMA_DESCRIPTION,
+                query_result="",
+                query_error=True,
+                feedback=f"Unknown task_id: {action.task_id}. Valid: task_1, task_2, task_3",
+                reward=0.0,
+                done=False,
+                step_count=self._step_count,
+                max_steps=MAX_STEPS_PER_TASK,
+            )
+
+        task = ALL_TASKS[action.task_id]
+        self._current_task = task
+
+        # Execute agent's SQL
+        agent_rows, agent_cols, error = execute_query(self._db, action.sql_query)
+
+        # Grade it
+        score, feedback = grade_query(task, agent_rows, agent_cols, error)
+
+        if score > self._best_reward:
+            self._best_reward = score
+
+        # Format result for observation
+        if error:
+            result_str = f"ERROR: {error}"
+            has_error = True
+        elif agent_rows is not None and agent_cols is not None:
+            header = " | ".join(agent_cols)
+            separator = "-" * len(header)
+            row_strs = [" | ".join(str(v) for v in row) for row in agent_rows[:20]]
+            result_str = f"{header}\n{separator}\n" + "\n".join(row_strs)
+            if len(agent_rows) > 20:
+                result_str += f"\n... ({len(agent_rows) - 20} more rows)"
+            has_error = False
+        else:
+            result_str = "(no result rows returned)"
+            has_error = False
+
+        # End episode if perfect score or max steps
+        is_done = False
+        if score >= 1.0:
+            is_done = True
+            feedback += " Task completed perfectly!"
+        elif self._step_count >= MAX_STEPS_PER_TASK:
+            is_done = True
+            feedback += f" Maximum steps ({MAX_STEPS_PER_TASK}) reached."
+
+        self._done = is_done
+
+        return SQLObservation(
+            task_id=task.task_id,
+            task_description=task.description,
+            difficulty=task.difficulty,
+            schema_description=SCHEMA_DESCRIPTION,
+            query_result=result_str,
+            query_error=has_error,
+            feedback=feedback,
+            reward=score,
+            done=is_done,
+            step_count=self._step_count,
+            max_steps=MAX_STEPS_PER_TASK,
+        )
+
+    @property
+    def state(self) -> State:
+        return self._state

server/tasks.py

@@ -0,0 +1,129 @@
+"""
+Task definitions for the SQL Query Environment.
+
+Each task has:
+  - task_id: unique identifier
+  - description: natural language question
+  - difficulty: easy / medium / hard
+  - expected_sql: a reference SQL that produces the correct answer
+  - expected_columns: the column names in the expected result
+  - expected_rows: the exact rows expected (deterministic)
+"""
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class TaskDefinition:
+    task_id: str
+    description: str
+    difficulty: str
+    expected_sql: str
+    expected_columns: list[str]
+    expected_rows: list[tuple]
+
+
+# ──────────────────────────────────────────────
+# TASK 1 — EASY: Single table, simple WHERE + ORDER BY
+# ──────────────────────────────────────────────
+TASK_1 = TaskDefinition(
+    task_id="task_1",
+    description=(
+        "List the names and salaries of all active employees who earn more "
+        "than 70000, ordered by salary descending."
+    ),
+    difficulty="easy",
+    expected_sql="""
+        SELECT name, salary
+        FROM employees
+        WHERE is_active = 1 AND salary > 70000
+        ORDER BY salary DESC;
+    """,
+    expected_columns=["name", "salary"],
+    expected_rows=[
+        ("Ravi Krishnan", 95000.0),
+        ("Priya Patel", 92000.0),
+        ("Amit Rao", 88000.0),
+        ("Arjun Sharma", 85000.0),
+        ("Lakshmi Iyer", 82000.0),
+        ("Rahul Verma", 78000.0),
+        ("Anita Desai", 72000.0),
+        ("Nikhil Agarwal", 71000.0),
+    ],
+)
+
+# ──────────────────────────────────────────────
+# TASK 2 — MEDIUM: JOIN + GROUP BY + HAVING
+# ──────────────────────────────────────────────
+TASK_2 = TaskDefinition(
+    task_id="task_2",
+    description=(
+        "For each department, show the department name, the number of active "
+        "employees, and the average salary of active employees. Only include "
+        "departments that have more than 2 active employees. "
+        "Order by average salary descending."
+    ),
+    difficulty="medium",
+    expected_sql="""
+        SELECT d.name, COUNT(e.id) AS num_employees, ROUND(AVG(e.salary), 2) AS avg_salary
+        FROM departments d
+        JOIN employees e ON d.id = e.department_id
+        WHERE e.is_active = 1
+        GROUP BY d.id, d.name
+        HAVING COUNT(e.id) > 2
+        ORDER BY avg_salary DESC;
+    """,
+    expected_columns=["name", "num_employees", "avg_salary"],
+    expected_rows=[
+        ("Engineering", 4, 87500.0),
+        ("Sales", 3, 70333.33),
+        ("Marketing", 3, 65666.67),
+    ],
+)
+
+# ──────────────────────────────────────────────
+# TASK 3 — HARD: Subquery + multiple JOINs + complex logic
+# ──────────────────────────────────────────────
+TASK_3 = TaskDefinition(
+    task_id="task_3",
+    description=(
+        "Find all active employees who lead at least one active project and "
+        "whose salary is above the overall average salary of all active employees. "
+        "For each such employee, show their name, salary, department name, and "
+        "the total budget of active projects they lead. "
+        "Order by total project budget descending."
+    ),
+    difficulty="hard",
+    expected_sql="""
+        SELECT
+            e.name,
+            e.salary,
+            d.name AS department_name,
+            SUM(p.budget) AS total_project_budget
+        FROM employees e
+        JOIN departments d ON e.department_id = d.id
+        JOIN projects p ON p.lead_employee_id = e.id
+        WHERE e.is_active = 1
+          AND p.status = 'active'
+          AND e.salary > (
+              SELECT AVG(salary) FROM employees WHERE is_active = 1
+          )
+        GROUP BY e.id, e.name, e.salary, d.name
+        ORDER BY total_project_budget DESC;
+    """,
+    expected_columns=["name", "salary", "department_name", "total_project_budget"],
+    expected_rows=[
+        ("Ravi Krishnan", 95000.0, "Engineering", 150000.0),
+        ("Priya Patel", 92000.0, "Engineering", 120000.0),
+        ("Arjun Sharma", 85000.0, "Engineering", 80000.0),
+    ],
+)
+
+# All tasks in order
+ALL_TASKS = {
+    "task_1": TASK_1,
+    "task_2": TASK_2,
+    "task_3": TASK_3,
+}
+
+TASK_LIST = [TASK_1, TASK_2, TASK_3]