agent setup

.gitignore

@@ -0,0 +1,15 @@
+
+# Security / Secrets
+.env
+
+# Virtual Environments
+venv/
+env/
+
+# Python Cache
+__pycache__/
+*.pyc
+
+# OS generated files
+.DS_Store
+Thumbs.db

README.md

@@ -7,66 +7,240 @@ sdk: docker
 pinned: false
 ---
 
-# SQL Analyst OpenEnv
+# 📊 SQL Analyst OpenEnv
 
-A real-world OpenEnv environment where an AI agent writes SQL queries against an e-commerce database to answer business questions.
+> A real-world OpenEnv environment where an AI agent must write correct SQL queries
+> against a live e-commerce database to answer business analytics questions.
+
+**Live Demo:** https://p-karthik-mohan-sql-analyst-env.hf.space/docs
+
+---
+
+## What Is This?
+
+This environment simulates the daily work of a **data analyst at an e-commerce company**.
+The AI agent receives a natural language business question, explores the database schema,
+and must produce a correct SQL query to answer it.
+
+Unlike toy environments, this is a task that real analysts perform every day —
+making it a meaningful benchmark for AI reasoning and code generation.
+
+---
 
 ## Tasks
 
-| ID | Difficulty | Description |
-|----|-----------|-------------|
-| 1  | Easy      | Count completed orders in 2024 |
-| 2  | Medium    | Top 5 customers by revenue (JOIN + GROUP BY) |
-| 3  | Hard      | Category revenue ranking (CTE + window function) |
+Three tasks of increasing difficulty, each graded by an automated SQL result comparator.
+
+### Task 1 — Easy
+**"How many completed orders were placed in 2024?"**
+- Requires: COUNT, WHERE, date filtering
+- Tests: basic aggregation and filtering
+- Expected output: single row, single column
+
+### Task 2 — Medium
+**"Find the top 5 customers by total revenue from completed orders."**
+- Requires: JOIN, GROUP BY, SUM, ORDER BY, LIMIT
+- Tests: multi-table joins and aggregation
+- Expected output: 5 rows with first_name, last_name, total_revenue
+
+### Task 3 — Hard
+**"Rank product categories by total revenue using a window function."**
+- Requires: CTE (WITH), JOIN, GROUP BY, RANK() OVER (...)
+- Tests: advanced SQL — CTEs and window functions
+- Expected output: all categories with total_revenue and revenue_rank
+
+---
 
 ## API Endpoints
 
-- `POST /reset` — load a task
-- `POST /step` — submit a SQL query, get reward 0.0–1.0
-- `GET /state` — current session state
-- `GET /docs` — interactive API documentation
+Base URL: https://p-karthik-mohan-sql-analyst-env.hf.space
 
-## Quick Start
+### POST /reset
+Load a task. Always call this first.
 
-```bash
-# Reset to task 1
-curl -X POST https://p-karthik-mohan-sql-analyst-env.hf.space/reset \
-  -H "Content-Type: application/json" \
-  -d '{"task_id": 1}'
-
-# Submit a SQL query
-curl -X POST https://p-karthik-mohan-sql-analyst-env.hf.space/step \
-  -H "Content-Type: application/json" \
-  -d '{"action": "SELECT COUNT(*) AS total_orders FROM orders WHERE status = 'completed'"}'
+Request:
+```json
+{"task_id": 1}
+```
+
+Response:
+```json
+{
+  "observation": {
+    "task_id": 1,
+    "difficulty": "easy",
+    "task_description": "Find the total number of completed orders...",
+    "schema": "Tables:\n  customers (...)\n  products (...)\n  orders (...)",
+    "hint": "Use COUNT with WHERE filters on status and order_date"
+  },
+  "info": {"message": "Task 1 loaded. Use POST /step with your SQL query."}
+}
 ```
 
+### POST /step
+Submit a SQL query. Returns reward and feedback.
+
+Request:
+```json
+{"action": "SELECT COUNT(*) AS total_orders FROM orders WHERE status = 'completed'"}
+```
+
+Response:
+```json
+{
+  "observation": {
+    "result_preview": [{"total_orders": 312}],
+    "reward_breakdown": {
+      "column_score": 0.30,
+      "row_score": 0.30,
+      "value_score": 0.40,
+      "total_reward": 1.0
+    }
+  },
+  "reward": 1.0,
+  "done": true
+}
+```
+
+### GET /state
+Get current session — task info, all attempts, best score.
+
+---
+
+## Observation Space
+
+| Field | Type | Description |
+|---|---|---|
+| task_description | string | Natural language business question |
+| schema | string | All table names, columns, types, row counts |
+| hint | string | Guidance on which SQL constructs to use |
+| result_preview | array | First 5 rows of the agent query result |
+| result_row_count | integer | Total rows returned by agent query |
+| reward_breakdown | object | Sub-scores for columns, rows, values |
+
+---
+
+## Action Space
+
+| Property | Value |
+|---|---|
+| Type | string |
+| Format | Valid SQLite SELECT or WITH statement |
+| Restrictions | No INSERT, UPDATE, DELETE, DROP |
+
+Example actions:
+```sql
+-- Easy
+SELECT COUNT(*) AS total_orders FROM orders WHERE status = 'completed'
+
+-- Medium
+SELECT c.first_name, c.last_name, SUM(o.total_amount) AS total_revenue
+FROM orders o JOIN customers c ON o.customer_id = c.customer_id
+WHERE o.status = 'completed'
+GROUP BY o.customer_id ORDER BY total_revenue DESC LIMIT 5
+
+-- Hard
+WITH rev AS (
+  SELECT p.category, SUM(o.total_amount) AS total_revenue
+  FROM orders o JOIN products p ON o.product_id = p.product_id
+  WHERE o.status = 'completed' GROUP BY p.category
+)
+SELECT category, total_revenue, RANK() OVER (ORDER BY total_revenue DESC) AS revenue_rank
+FROM rev ORDER BY revenue_rank ASC
+```
+
+---
+
 ## Reward Function
 
-Partial credit scoring — 0.0 to 1.0:
+Partial credit score from 0.0 to 1.0 with three components:
 
-- **0.30** — correct column names
-- **0.30** — correct row count
-- **0.40** — correct cell values
+| Component | Weight | How it is measured |
+|---|---|---|
+| Column names | 0.30 | Fraction of expected column names present |
+| Row count | 0.30 | Ratio of returned rows vs expected rows |
+| Cell values | 0.40 | Fraction of cells matching expected values |
 
-## Running inference.py
+Even an imperfect query receives meaningful feedback — not just pass/fail.
+This gives the agent a gradient signal to improve from.
+
+---
+
+## Database Schema
+
+A realistic e-commerce dataset with 600 orders, 100 customers, 30 products.
+
+```
+customers  (customer_id, first_name, last_name, email, city, signup_date)
+products   (product_id, product_name, category, price, stock)
+orders     (order_id, customer_id, product_id, quantity, total_amount, order_date, status)
+```
+
+All orders are dated in 2024. Status values: completed, pending, cancelled.
+
+---
+
+## Running the Baseline Agent
 
 ```bash
+pip install openai requests
+
 export API_BASE_URL=https://api.openai.com/v1
 export MODEL_NAME=gpt-4o-mini
 export HF_TOKEN=your_api_key_here
+
 python inference.py
 ```
 
-## Database Schema
+Expected output:
+```
+Tasks solved   : 3 / 3
+Average reward : 1.000 / 1.000
+Results saved to results.json
+```
 
+---
+
+## Local Setup
+
+```bash
+git clone https://huggingface.co/spaces/P-Karthik-Mohan/sql-analyst-env
+cd sql-analyst-env
+pip install -r requirements.txt
+python seed.py
+uvicorn main:app --host 0.0.0.0 --port 7860
 ```
-customers  (customer_id, first_name, last_name, email, city, signup_date)  — 100 rows
-products   (product_id, product_name, category, price, stock)              — 30 rows
-orders     (order_id, customer_id, product_id, quantity, total_amount, order_date, status) — 600 rows
+
+---
+
+## Docker
+
+```bash
+docker build -t sql-analyst-env .
+docker run -p 7860:7860 sql-analyst-env
 ```
 
+---
+
 ## Hardware Requirements
 
-- 2 vCPU, 8GB RAM minimum
-- No GPU required
-- SQLite — no external database needed
+- CPU: 1-2 vCPU
+- RAM: 8GB
+- GPU: Not required
+- Inference time: Under 5 minutes for all 3 tasks
+
+---
+
+## Project Structure
+
+```
+sql-analyst-env/
+├── main.py          # FastAPI server
+├── seed.py          # Database seeder
+├── inference.py     # Baseline AI agent
+├── openenv.yaml     # OpenEnv specification
+├── Dockerfile       # Container for HF Spaces
+├── requirements.txt # Python dependencies
+└── data/
+    └── ecommerce.db # SQLite database
+```

inference.py

@@ -3,32 +3,23 @@ inference.py — Baseline AI agent for SQL Analyst OpenEnv
 ---------------------------------------------------------
 Uses the OpenAI client (pointed at any compatible LLM via API_BASE_URL)
 to solve all 3 tasks by interacting with the running FastAPI environment.
-
-Environment variables required:
-  API_BASE_URL  — LLM API base URL  (e.g. https://api.openai.com/v1)
-  MODEL_NAME    — model to use       (e.g. gpt-4o-mini)
-  HF_TOKEN      — Hugging Face token (used as the API key)
-
-Usage:
-  python inference.py
 """
-
 import os
 import sys
-import json
 import time
 import requests
 from openai import OpenAI
-
+from dotenv import load_dotenv
+load_dotenv()
 # ── Configuration ─────────────────────────────────────────────────────────────
 
-ENV_BASE_URL = "http://127.0.0.1:7860"   # where FastAPI server is running
+ENV_BASE_URL = "https://p-karthik-mohan-sql-analyst-env.hf.space"   # live HF Space
 MAX_ATTEMPTS = 5                         # max SQL attempts per task
-TASK_IDS     = [1, 2, 3]                # tasks to solve
+TASK_IDS     = [1, 2, 3]                 # tasks to solve
 
 API_BASE_URL = "https://api.groq.com/openai/v1"
 MODEL_NAME   = "llama-3.1-8b-instant"
-HF_TOKEN     = "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"
+HF_TOKEN     = os.environ.get("HF_TOKEN")
 
 # ── OpenAI Client ─────────────────────────────────────────────────────────────
 
@@ -44,33 +35,29 @@ def env_reset(task_id: int) -> dict:
     r.raise_for_status()
     return r.json()
 
-
 def env_step(sql: str) -> dict:
     r = requests.post(f"{ENV_BASE_URL}/step", json={"action": sql})
     r.raise_for_status()
     return r.json()
 
-
 def env_state() -> dict:
     r = requests.get(f"{ENV_BASE_URL}/state")
     r.raise_for_status()
     return r.json()
 
-
 def wait_for_server(retries: int = 10, delay: float = 2.0):
     """Wait until the FastAPI server is ready."""
     print("Waiting for environment server...")
     for i in range(retries):
         try:
-            r = requests.get(f"{ENV_BASE_URL}/health", timeout=3)
+            r = requests.get(f"{ENV_BASE_URL}/docs")
             if r.status_code == 200:
-                print("Server is ready.\n")
+                print("Server is up!")
                 return
-        except requests.exceptions.ConnectionError:
+        except requests.ConnectionError:
             pass
-        print(f"  Not ready yet, retrying in {delay}s... ({i+1}/{retries})")
         time.sleep(delay)
-    print("ERROR: Server did not start in time. Is uvicorn running?")
+    print("ERROR: Server did not start in time.")
     sys.exit(1)
 
 # ── LLM SQL Generator ─────────────────────────────────────────────────────────
@@ -86,7 +73,6 @@ Rules:
 - If a previous attempt scored less than 1.0, study the feedback and fix the query.
 """
 
-
 def build_user_prompt(
     task_description: str,
     schema: str,
@@ -105,22 +91,13 @@ Hint: {hint}
 Attempt number: {attempt}
 """
     if previous_attempts:
-        prompt += "\nYour previous attempts and their scores:\n"
-        for prev in previous_attempts[-3:]:   # show last 3 only
-            prompt += f"""
-  Attempt {prev['attempt']}:
-    SQL:    {prev['sql']}
-    Reward: {prev['reward']} / 1.0
-    Columns expected : {prev['details'].get('expected_columns', [])}
-    Columns you gave : {prev['details'].get('agent_columns', [])}
-    Rows expected    : {prev['details'].get('expected_row_count', '?')}
-    Rows you gave    : {prev['details'].get('agent_row_count', '?')}
-    Sub-scores       : columns={prev['details'].get('column_score', 0):.2f}  rows={prev['details'].get('row_score', 0):.2f}  values={prev['details'].get('value_score', 0):.2f}
-"""
+        prompt += "\nPrevious attempts:\n"
+        for i, prev in enumerate(previous_attempts):
+            prompt += f"--- Attempt {i+1} ---\nSQL: {prev['sql']}\nReward: {prev['reward']}\n\n"
+        
     prompt += "\nWrite the corrected SQL query now:"
     return prompt
 
-
 def ask_llm(task_description: str, schema: str, hint: str,
             attempt: int, previous_attempts: list) -> str:
     """Call the LLM and return a SQL string."""
@@ -133,18 +110,14 @@ def ask_llm(task_description: str, schema: str, hint: str,
     response = client.chat.completions.create(
         model=MODEL_NAME,
         messages=messages,
-        temperature=0.0,       # deterministic — we want correct SQL, not creative SQL
+        temperature=0.0,
         max_tokens=512,
     )
     sql = response.choices[0].message.content.strip()
 
     # Strip markdown fences if model wraps in ```sql ... ```
     if sql.startswith("```"):
-        lines = sql.split("\n")
-        sql = "\n".join(
-            line for line in lines
-            if not line.strip().startswith("```")
-        ).strip()
+        sql = sql.split("\n", 1)[-1].rsplit("\n", 1)[0].replace("```", "").strip()
 
     return sql
 
@@ -172,47 +145,28 @@ def solve_task(task_id: int) -> dict:
     final_sql   = ""
 
     for attempt in range(1, MAX_ATTEMPTS + 1):
-        print(f"  Attempt {attempt}/{MAX_ATTEMPTS} — asking LLM...")
-
-        # Get SQL from LLM
+        print(f"Attempt {attempt}/{MAX_ATTEMPTS}...")
         sql = ask_llm(task_desc, schema, hint, attempt, previous_attempts)
-        print(f"  SQL: {sql[:120]}{'...' if len(sql) > 120 else ''}")
-
-        # Submit to environment
+        print(f"Generated SQL: {sql}")
+        
         step_resp = env_step(sql)
-        reward    = step_resp["reward"]
-        done      = step_resp["done"]
-        details   = step_resp["observation"].get("reward_breakdown", {})
-
-        print(f"  Reward: {reward:.3f}  "
-              f"(cols={details.get('column_score',0):.2f}  "
-              f"rows={details.get('row_score',0):.2f}  "
-              f"vals={details.get('value_score',0):.2f})")
-
+        reward = step_resp["reward"]
+        print(f"Reward: {reward}")
+        
         best_reward = max(best_reward, reward)
-        final_sql   = sql
-
-        # Store attempt for LLM feedback
-        previous_attempts.append({
-            "attempt": attempt,
-            "sql":     sql,
-            "reward":  reward,
-            "details": details,
-        })
-
-        if done:
-            print(f"  PERFECT SCORE on attempt {attempt}!")
+        final_sql = sql
+        
+        if reward >= 1.0:
+            print("Task solved successfully!")
             break
-        elif reward >= 0.8:
-            print(f"  Score is close ({reward:.3f}). Trying to improve...")
-        else:
-            print(f"  Score is low ({reward:.3f}). Refining query...")
+            
+        previous_attempts.append({"sql": sql, "reward": reward})
 
     return {
         "task_id":      task_id,
         "difficulty":   difficulty,
         "best_reward":  best_reward,
-        "attempts":     len(previous_attempts),
+        "attempts":     len(previous_attempts) + (1 if best_reward >= 1.0 else 0),
         "final_sql":    final_sql,
         "solved":       best_reward >= 1.0,
     }
@@ -228,8 +182,8 @@ def main():
 
     results = []
     for task_id in TASK_IDS:
-        result = solve_task(task_id)
-        results.append(result)
+        res = solve_task(task_id)
+        results.append(res)
 
     # ── Final Summary ─────────────────────────────────────────────────────────
     print(f"\n{'='*60}")
@@ -237,27 +191,16 @@ def main():
     print('='*60)
 
     total_score = 0.0
+    solved_count = 0
     for r in results:
-        status = "SOLVED" if r["solved"] else f"best={r['best_reward']:.3f}"
-        print(f"  Task {r['task_id']} ({r['difficulty']:6s})  {status}  "
-              f"in {r['attempts']} attempt(s)")
         total_score += r["best_reward"]
+        if r["solved"]:
+            solved_count += 1
+        print(f"Task {r['task_id']} ({r['difficulty']}): Reward = {r['best_reward']:.2f} | Solved = {r['solved']}")
 
-    avg_score = total_score / len(results)
-    print(f"\n  Average reward : {avg_score:.3f} / 1.000")
-    print(f"  Tasks solved   : {sum(1 for r in results if r['solved'])} / {len(results)}")
-
-    # Save results to file (useful for judges / CI)
-    output_path = "results.json"
-    with open(output_path, "w") as f:
-        json.dump({
-            "results":     results,
-            "avg_score":   round(avg_score, 3),
-            "tasks_solved": sum(1 for r in results if r["solved"]),
-        }, f, indent=2)
-    print(f"\n  Results saved to {output_path}")
-
-    return avg_score
+    avg_score = total_score / len(results) if results else 0
+    print(f"\nTasks solved   : {solved_count} / {len(results)}")
+    print(f"Average reward : {avg_score:.3f} / 1.000")
 
 
 if __name__ == "__main__":

main.py

@@ -4,128 +4,119 @@ import json
 import re
 from datetime import datetime
 from typing import Any, Optional
+import uuid
 
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+import gradio as gr
+from ui import build_ui
 
 # ── Config ────────────────────────────────────────────────────────────────────
 DB_PATH = os.path.join("data", "ecommerce.db")
-app = FastAPI(title="SQL Analyst OpenEnv", version="1.0.0")
+# Under /api for direct agent access, root for UI
+api_app = FastAPI(
+    title="SQL Analyst OpenEnv API", 
+    version="1.1.0", 
+    docs_url="/docs", 
+    redoc_url="/redoc",
+    description="A real-world benchmark environment for AI agents writing SQL. Use `/reset` to select a task, `/step` to submit a query and receive partial-credit feedback, and `/state` to track history."
+)
+
 
 # ── Pydantic Models ───────────────────────────────────────────────────────────
 
 class StepRequest(BaseModel):
-    action: str          # The SQL query the agent submits
+    session_id: str = Field("default", description="Unique session ID to prevent collisions.")
+    action: str = Field(..., description="A valid SQLite SELECT or WITH statement.", examples=["SELECT COUNT(*) AS total_orders FROM orders;"])
 
 class StepResponse(BaseModel):
-    observation: dict
-    reward: float        # 0.0 – 1.0
-    done: bool
-    info: dict
+    observation: dict = Field(description="Contains result_preview (first 5 rows), reward_breakdown, and any sql execution errors.")
+    reward: float = Field(description="Partial credit reward from 0.0 (wrong) to 1.0 (perfect).")
+    done: bool = Field(description="True if reward == 1.0 (Task solved).")
+    info: dict = Field(description="System message and attempt count.")
 
 class ResetRequest(BaseModel):
-    task_id: int         # 1 = Easy, 2 = Medium, 3 = Hard
+    task_id: int = Field(..., description="ID of the task to load (1 to 5).", examples=[1])
+    session_id: str = Field("default", description="Unique session ID.")
 
 class ResetResponse(BaseModel):
-    observation: dict
-    info: dict
+    observation: dict = Field(description="Contains task_description, schema, and hint to be passed to the agent.")
+    info: dict = Field(description="System message.")
+
+class StateRequest(BaseModel):
+    session_id: str = Field("default", description="Unique session ID.")
 
 class StateResponse(BaseModel):
-    task_id: int
-    task_description: str
-    schema_info: str
-    attempts: int
-    best_reward: float
-    history: list
+    session_id: str
+    task_id: Optional[int] = Field(description="Currently active Task ID.")
+    task_description: Optional[str]
+    schema_info: str = Field(description="Raw string representation of DB Schema.")
+    attempts: int = Field(description="Number of queries submitted so far on active task.")
+    best_reward: float = Field(description="Highest partial credit achieved.")
+    history: list = Field(description="Log of all queries executed.")
 
 # ── Task Definitions ──────────────────────────────────────────────────────────
 
 TASKS = {
     1: {
-        "description": (
-            "Find the total number of completed orders placed in the year 2024. "
-            "Return a single number with column name: total_orders"
-        ),
+        "description": "Find the total number of completed orders placed in the year 2024. Return a single number with column name: total_orders",
         "difficulty": "easy",
         "hint": "Use COUNT with WHERE filters on status and order_date",
-        "answer_query": """
-            SELECT COUNT(*) AS total_orders
-            FROM orders
-            WHERE status = 'completed'
-              AND order_date LIKE '2024%'
-        """,
+        "answer_query": "SELECT COUNT(*) AS total_orders FROM orders WHERE status = 'completed' AND order_date LIKE '2024%'",
     },
     2: {
-        "description": (
-            "Find the top 5 customers by total revenue (sum of total_amount for completed orders only). "
-            "Return columns: first_name, last_name, total_revenue. "
-            "Order by total_revenue descending."
-        ),
+        "description": "Find the top 5 customers by total revenue (sum of total_amount for completed orders only). Return columns: first_name, last_name, total_revenue. Order by total_revenue descending.",
         "difficulty": "medium",
         "hint": "JOIN orders with customers, GROUP BY customer, filter completed, ORDER and LIMIT",
-        "answer_query": """
-            SELECT c.first_name, c.last_name,
-                   ROUND(SUM(o.total_amount), 2) AS total_revenue
-            FROM orders o
-            JOIN customers c ON o.customer_id = c.customer_id
-            WHERE o.status = 'completed'
-            GROUP BY o.customer_id
-            ORDER BY total_revenue DESC
-            LIMIT 5
-        """,
+        "answer_query": "SELECT c.first_name, c.last_name, ROUND(SUM(o.total_amount), 2) AS total_revenue FROM orders o JOIN customers c ON o.customer_id = c.customer_id WHERE o.status = 'completed' GROUP BY o.customer_id ORDER BY total_revenue DESC LIMIT 5",
     },
     3: {
-        "description": (
-            "For each product category, calculate the total revenue (completed orders only) "
-            "and rank categories by revenue using a window function. "
-            "Return columns: category, total_revenue, revenue_rank. "
-            "Order by revenue_rank ascending."
-        ),
+        "description": "For each product category, calculate the total revenue (completed orders only) and rank categories by revenue using a window function. Return columns: category, total_revenue, revenue_rank. Order by revenue_rank ascending.",
         "difficulty": "hard",
         "hint": "Use SUM with GROUP BY inside a CTE, then apply RANK() OVER (ORDER BY ...) on the result",
-        "answer_query": """
-            WITH category_revenue AS (
-                SELECT p.category,
-                       SUM(o.total_amount) AS total_revenue
-                FROM orders o
-                JOIN products p ON o.product_id = p.product_id
-                WHERE o.status = 'completed'
-                GROUP BY p.category
-            )
-            SELECT category,
-                   total_revenue,
-                   RANK() OVER (ORDER BY total_revenue DESC) AS revenue_rank
-            FROM category_revenue
-            ORDER BY revenue_rank ASC
-        """,
+        "answer_query": "WITH category_revenue AS ( SELECT p.category, SUM(o.total_amount) AS total_revenue FROM orders o JOIN products p ON o.product_id = p.product_id WHERE o.status = 'completed' GROUP BY p.category ) SELECT category, total_revenue, RANK() OVER (ORDER BY total_revenue DESC) AS revenue_rank FROM category_revenue ORDER BY revenue_rank ASC",
+    },
+    4: {
+        "description": "Find the average price of products in each category, but only for categories that have more than 2 products. Return columns: category, avg_price.",
+        "difficulty": "medium",
+        "hint": "Use GROUP BY with HAVING COUNT(...) > 2.",
+        "answer_query": "SELECT category, ROUND(AVG(price), 2) AS avg_price FROM products GROUP BY category HAVING COUNT(product_id) > 2",
+    },
+    5: {
+        "description": "Identify customers who have ordered products from both the 'Electronics' and 'Clothing' categories. Return columns: customer_id, first_name.",
+        "difficulty": "hard",
+        "hint": "Use INTERSECT on two queries, or GROUP BY customer HAVING COUNT(DISTINCT category) = 2.",
+        "answer_query": "SELECT DISTINCT c.customer_id, c.first_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN products p ON o.product_id = p.product_id WHERE p.category = 'Electronics' INTERSECT SELECT DISTINCT c.customer_id, c.first_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN products p ON o.product_id = p.product_id WHERE p.category = 'Clothing'",
     },
 }
 
-# ── In-Memory Session State ───────────────────────────────────────────────────
+# ── Sessions ──────────────────────────────────────────────────────────────────
 
-session = {
-    "task_id": None,
-    "task": None,
-    "expected_rows": None,
-    "expected_columns": None,
-    "attempts": 0,
-    "best_reward": 0.0,
-    "history": [],
-}
+sessions = {}
+
+def get_session(sid: str):
+    if sid not in sessions:
+        sessions[sid] = {
+            "task_id": None, "task": None, 
+            "expected_rows": None, "expected_columns": None,
+            "attempts": 0, "best_reward": 0.0, "history": []
+        }
+    return sessions[sid]
 
 # ── Database Helpers ──────────────────────────────────────────────────────────
 
+def progress_handler():
+    raise sqlite3.OperationalError("Query execution aborted: Timed out or exceeded instruction limits. Hint: Too complex CROSS JOIN?")
+
 def get_connection():
     if not os.path.exists(DB_PATH):
-        raise HTTPException(
-            status_code=500,
-            detail=f"Database not found at {DB_PATH}. Run seed.py first."
-        )
+        raise HTTPException(status_code=500, detail=f"Database not found at {DB_PATH}.")
     conn = sqlite3.connect(DB_PATH)
     conn.row_factory = sqlite3.Row
+    # Security feature: Prevent DOS
+    conn.set_progress_handler(progress_handler, 500000)
     return conn
 
-
 def get_schema_info() -> str:
     conn = get_connection()
     cur = conn.cursor()
@@ -142,9 +133,7 @@ def get_schema_info() -> str:
     conn.close()
     return "Tables:\n" + "\n".join(schema_parts)
 
-
 def run_query(sql: str) -> tuple[list[dict], list[str]]:
-    """Run a SQL query and return (rows_as_dicts, column_names)."""
     conn = get_connection()
     cur = conn.cursor()
     cur.execute(sql)
@@ -153,9 +142,8 @@ def run_query(sql: str) -> tuple[list[dict], list[str]]:
     conn.close()
     return rows, columns
 
-
-def compute_expected():
-    """Run the answer query and cache the expected result."""
+def compute_expected(sid: str):
+    session = get_session(sid)
     task = session["task"]
     rows, columns = run_query(task["answer_query"])
     session["expected_rows"] = rows
@@ -163,29 +151,18 @@ def compute_expected():
 
 # ── Reward Function ───────────────────────────────────────────────────────────
 
-def compute_reward(agent_rows: list[dict], agent_cols: list[str]) -> tuple[float, dict]:
-    """
-    Partial scoring reward function — 0.0 to 1.0.
-
-    Breakdown:
-      - 0.30  correct column names
-      - 0.30  correct number of rows
-      - 0.40  correct values (cell-level match)
-    """
+def compute_reward(sid: str, agent_rows: list[dict], agent_cols: list[str]) -> tuple[float, dict]:
+    session = get_session(sid)
     expected_rows = session["expected_rows"]
     expected_cols = session["expected_columns"]
     details = {}
 
-    # ── Column score (0.30) ──────────────────────────────────────────────────
     agent_cols_lower   = [c.lower() for c in agent_cols]
     expected_cols_lower = [c.lower() for c in expected_cols]
     col_matches = sum(1 for c in expected_cols_lower if c in agent_cols_lower)
     col_score = (col_matches / len(expected_cols_lower)) * 0.30 if expected_cols_lower else 0.0
     details["column_score"] = round(col_score, 3)
-    details["expected_columns"] = expected_cols
-    details["agent_columns"] = agent_cols
 
-    # ── Row count score (0.30) ───────────────────────────────────────────────
     expected_count = len(expected_rows)
     agent_count    = len(agent_rows)
     if expected_count == 0:
@@ -194,21 +171,14 @@ def compute_reward(agent_rows: list[dict], agent_cols: list[str]) -> tuple[float
         row_ratio = min(agent_count, expected_count) / max(agent_count, expected_count)
         row_score = row_ratio * 0.30
     details["row_score"] = round(row_score, 3)
-    details["expected_row_count"] = expected_count
-    details["agent_row_count"]    = agent_count
 
-    # ── Value match score (0.40) ─────────────────────────────────────────────
     if not expected_rows or not agent_rows:
         value_score = 0.0
     else:
         def normalize(v):
-            if v is None:
-                return ""
-            try:
-                return str(round(float(v), 1))
-            except (ValueError, TypeError):
-                return str(v).strip().lower()
-
+            if v is None: return ""
+            try: return str(round(float(v), 1))
+            except: return str(v).strip().lower()
 
         matched_cells = 0
         total_cells   = len(expected_rows) * len(expected_cols_lower)
@@ -216,10 +186,8 @@ def compute_reward(agent_rows: list[dict], agent_cols: list[str]) -> tuple[float
         for exp_row, agt_row in zip(expected_rows, agent_rows):
             for col in expected_cols_lower:
                 exp_val = normalize(exp_row.get(col) or exp_row.get(col.upper()))
-                # Try matching by column name first, then by position
                 agt_val = normalize(agt_row.get(col) or agt_row.get(col.upper()))
                 if not agt_val:
-                    # Fall back to positional match
                     exp_idx = expected_cols_lower.index(col)
                     if exp_idx < len(agent_cols):
                         pos_col = agent_cols[exp_idx]
@@ -236,132 +204,80 @@ def compute_reward(agent_rows: list[dict], agent_cols: list[str]) -> tuple[float
 
 # ── Endpoints ─────────────────────────────────────────────────────────────────
 
-@app.post("/reset", response_model=ResetResponse)
+@api_app.post("/reset", response_model=ResetResponse, tags=["Agent Environment"], summary="Load Task", description="Initializes the environment with a random task, or resets the current task, returning the description, hints, and database schema.")
 def reset(req: ResetRequest):
-    """Start a new task. Call this before /step."""
-    if req.task_id not in TASKS:
-        raise HTTPException(status_code=400, detail="task_id must be 1, 2, or 3")
-
+    if req.task_id not in TASKS: raise HTTPException(status_code=400, detail="Invalid task_id")
+    session = get_session(req.session_id)
     session["task_id"]    = req.task_id
     session["task"]       = TASKS[req.task_id]
     session["attempts"]   = 0
     session["best_reward"] = 0.0
     session["history"]    = []
-
-    compute_expected()
-
-    schema = get_schema_info()
+    compute_expected(req.session_id)
     observation = {
         "task_id":          req.task_id,
         "difficulty":       session["task"]["difficulty"],
         "task_description": session["task"]["description"],
-        "schema":           schema,
+        "schema":           get_schema_info(),
         "hint":             session["task"]["hint"],
     }
-    return ResetResponse(
-        observation=observation,
-        info={"message": f"Task {req.task_id} loaded. Use POST /step with your SQL query."}
-    )
-
+    return ResetResponse(observation=observation, info={"message": f"Task {req.task_id} loaded."})
 
-@app.post("/step", response_model=StepResponse)
+@api_app.post("/step", response_model=StepResponse, tags=["Agent Environment"], summary="Execute SQL", description="Executes the valid SQLite `action` against the ecommerce database, evaluating the correctness using a partial 0.0-1.0 Reward function.")
 def step(req: StepRequest):
-    """Submit a SQL query. Returns reward 0.0–1.0 and feedback."""
-    if session["task_id"] is None:
-        raise HTTPException(status_code=400, detail="Call /reset first to load a task.")
-
+    session = get_session(req.session_id)
+    if session["task_id"] is None: raise HTTPException(status_code=400, detail="Call /reset first.")
     session["attempts"] += 1
     sql = req.action.strip()
 
-    # ── Safety: only allow SELECT statements ─────────────────────────────────
     if not re.match(r"^\s*(SELECT|WITH)\b", sql, re.IGNORECASE):
         return StepResponse(
-            observation={"error": "Only SELECT or WITH (CTE) statements are allowed."},
-            reward=0.0,
-            done=False,
-            info={"attempt": session["attempts"], "message": "Rejected: not a SELECT/WITH query."}
+            observation={"error": "Only SELECT or WITH allowed."}, reward=0.0, done=False,
+            info={"attempt": session["attempts"], "message": "Rejected"}
         )
 
-    # ── Run the agent's query ─────────────────────────────────────────────────
     try:
         agent_rows, agent_cols = run_query(sql)
     except Exception as e:
-        entry = {
-            "attempt": session["attempts"],
-            "sql": sql,
-            "reward": 0.0,
-            "error": str(e),
-        }
-        session["history"].append(entry)
+        session["history"].append({"attempt": session["attempts"], "sql": sql, "reward": 0.0, "error": str(e)})
         return StepResponse(
-            observation={"error": str(e), "sql_submitted": sql},
-            reward=0.0,
-            done=False,
-            info={"attempt": session["attempts"], "message": "SQL execution error."}
+            observation={"error": str(e), "sql_submitted": sql}, reward=0.0, done=False,
+            info={"attempt": session["attempts"], "message": "SQL Error"}
         )
 
-    # ── Score it ──────────────────────────────────────────────────────────────
-    reward, details = compute_reward(agent_rows, agent_cols)
+    reward, details = compute_reward(req.session_id, agent_rows, agent_cols)
     session["best_reward"] = max(session["best_reward"], reward)
-
     done = reward >= 1.0
-
-    entry = {
-        "attempt":   session["attempts"],
-        "sql":       sql,
-        "reward":    reward,
-        "details":   details,
-        "timestamp": datetime.now().isoformat(),
-    }
-    session["history"].append(entry)
+    session["history"].append({"attempt": session["attempts"], "sql": sql, "reward": reward, "details": details})
 
     observation = {
-        "task_id":          session["task_id"],
-        "task_description": session["task"]["description"],
-        "sql_submitted":    sql,
-        "result_preview":   agent_rows[:5],   # show first 5 rows
-        "result_row_count": len(agent_rows),
+        "task_id": session["task_id"], "task_description": session["task"]["description"],
+        "sql_submitted": sql, "result_preview": agent_rows[:5], "result_row_count": len(agent_rows),
         "reward_breakdown": details,
     }
-
     return StepResponse(
-        observation=observation,
-        reward=reward,
-        done=done,
-        info={
-            "attempt":      session["attempts"],
-            "best_reward":  session["best_reward"],
-            "message":      "Perfect score! Task complete." if done else "Keep refining your query.",
-        }
+        observation=observation, reward=reward, done=done,
+        info={"attempt": session["attempts"], "best_reward": session["best_reward"]}
     )
 
-
-@app.get("/state", response_model=StateResponse)
-def state():
-    """Get current session state — task info, attempts, history."""
-    if session["task_id"] is None:
-        raise HTTPException(status_code=400, detail="No active task. Call /reset first.")
-
+@api_app.get("/state", response_model=StateResponse, tags=["Diagnostics"], summary="Get Current State", description="Returns attempts, rewards, parsed tasks, and historical executed sql queries array.")
+def state(req: StateRequest):
+    session = get_session(req.session_id)
     return StateResponse(
+        session_id=req.session_id,
         task_id=session["task_id"],
-        task_description=session["task"]["description"],
+        task_description=session["task"]["description"] if session["task"] else None,
         schema_info=get_schema_info(),
         attempts=session["attempts"],
         best_reward=session["best_reward"],
         history=session["history"],
     )
 
-
-@app.get("/")
+@api_app.get("/")
 def root():
-    return {
-        "name":    "SQL Analyst OpenEnv",
-        "version": "1.0.0",
-        "tasks":   {k: {"difficulty": v["difficulty"], "description": v["description"]} for k, v in TASKS.items()},
-        "endpoints": ["/reset", "/step", "/state"],
-    }
+    return {"message": "API running at /api. Try the UI at root (handled by wrapper)!"}
 
+# ── Server Setup ──────────────────────────────────────────────────────────────
 
-@app.get("/health")
-def health():
-    return {"status": "ok", "db_exists": os.path.exists(DB_PATH)}
+demo = build_ui()
+app = gr.mount_gradio_app(api_app, demo, path="/")

make_awesome.py

@@ -0,0 +1,398 @@
+import os
+import re
+
+# 1. Update requirements.txt
+req_txt = "requirements.txt"
+with open(req_txt, "r") as f:
+    reqs = f.read()
+if "gradio" not in reqs:
+    with open(req_txt, "a") as f:
+        f.write("\ngradio==4.44.0\npandas\n")
+
+# 2. Create the Gradio UI file (ui.py)
+ui_code = """
+import gradio as gr
+import requests
+import pandas as pd
+import json
+import uuid
+
+ENV_BASE_URL = "http://127.0.0.1:7860"
+
+def new_session():
+    return str(uuid.uuid4())
+
+def load_task(task_id, sid):
+    try:
+        task_num = int(task_id.split()[1])
+        r = requests.post(f"{ENV_BASE_URL}/api/reset", json={"task_id": task_num, "session_id": sid})
+        if r.status_code != 200:
+            return f"Error: {r.text}", "", "", pd.DataFrame(), f"Error loading task {task_num}"
+            
+        data = r.json()
+        obs = data["observation"]
+        return obs["task_description"], obs["schema"], obs["hint"], pd.DataFrame(), "Task loaded. Write SQL below."
+    except Exception as e:
+        return str(e), "", "", pd.DataFrame(), "Error loading task"
+
+def run_sql(sql, sid):
+    if not sql.strip():
+        return pd.DataFrame(), "Please enter a SQL query.", "Error"
+    try:
+        r = requests.post(f"{ENV_BASE_URL}/api/step", json={"action": sql, "session_id": sid})
+        data = r.json()
+        obs = data.get("observation", {})
+        reward = data.get("reward", 0.0)
+        done = data.get("done", False)
+        
+        df = pd.DataFrame(obs.get("result_preview", []))
+        breakdown = obs.get("reward_breakdown", {})
+        
+        feedback = f"🎯 Reward: {reward:.2f} / 1.0\\n"
+        if breakdown:
+            feedback += f"Columns: {breakdown.get('column_score',0):.2f}, Rows: {breakdown.get('row_score',0):.2f}, Values: {breakdown.get('value_score',0):.2f}"
+        if "error" in obs:
+            feedback += f"\\n\\n⚠️ Error: {obs['error']}"
+        
+        return df, feedback, "✅ SOLVED!" if done else "Keep trying!"
+    except Exception as e:
+         return pd.DataFrame(), str(e), "Error"
+
+def build_ui():
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
+        gr.Markdown("# 📊 SQL Analyst OpenEnv - Hackathon Edition")
+        gr.Markdown("Test Human or AI performance on realistic E-Commerce SQL Data tasks. [API served at `/api`]")
+        
+        sid = gr.State(new_session)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                task_dropdown = gr.Dropdown(choices=["Task 1 (Easy)", "Task 2 (Medium)", "Task 3 (Hard)", "Task 4 (Medium)", "Task 5 (Hard)"], value="Task 1 (Easy)", label="Select Task")
+                btn_load = gr.Button("🔄 Load Task")
+                
+                desc = gr.Textbox(label="Business Question", interactive=False, lines=2)
+                hint = gr.Textbox(label="Hint", interactive=False)
+                schema = gr.Code(label="Database Schema", language="sql", interactive=False)
+            
+            with gr.Column(scale=2):
+                sql_input = gr.Code(label="SQL Editor", language="sql", lines=10)
+                btn_run = gr.Button("🚀 Run SQL", variant="primary")
+                
+                status_out = gr.Markdown("Ready.")
+                feedback_out = gr.Textbox(label="Feedback & Score", interactive=False)
+                grid_out = gr.Dataframe(label="Result Preview (First 5 Rows)")
+
+        btn_load.click(load_task, inputs=[task_dropdown, sid], outputs=[desc, schema, hint, grid_out, feedback_out])
+        btn_run.click(run_sql, inputs=[sql_input, sid], outputs=[grid_out, feedback_out, status_out])
+        
+    return demo
+"""
+with open("ui.py", "w", encoding="utf-8") as f:
+    f.write(ui_code.strip() + "\n")
+
+# 3. Rewrite main.py (adding concurrency fixes, security, and mounting gradio)
+main_py_code = """
+import sqlite3
+import os
+import json
+import re
+from datetime import datetime
+from typing import Any, Optional
+import uuid
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+import gradio as gr
+from ui import build_ui
+
+# ── Config ────────────────────────────────────────────────────────────────────
+DB_PATH = os.path.join("data", "ecommerce.db")
+# Under /api for direct agent access, root for UI
+api_app = FastAPI(title="SQL Analyst OpenEnv API", version="1.1.0")
+
+# ── Pydantic Models ───────────────────────────────────────────────────────────
+
+class StepRequest(BaseModel):
+    session_id: str = "default"
+    action: str
+
+class StepResponse(BaseModel):
+    observation: dict
+    reward: float
+    done: bool
+    info: dict
+
+class ResetRequest(BaseModel):
+    task_id: int
+    session_id: str = "default"
+
+class ResetResponse(BaseModel):
+    observation: dict
+    info: dict
+
+class StateRequest(BaseModel):
+    session_id: str = "default"
+
+class StateResponse(BaseModel):
+    session_id: str
+    task_id: Optional[int]
+    task_description: Optional[str]
+    schema_info: str
+    attempts: int
+    best_reward: float
+    history: list
+
+# ── Task Definitions ──────────────────────────────────────────────────────────
+
+TASKS = {
+    1: {
+        "description": "Find the total number of completed orders placed in the year 2024. Return a single number with column name: total_orders",
+        "difficulty": "easy",
+        "hint": "Use COUNT with WHERE filters on status and order_date",
+        "answer_query": "SELECT COUNT(*) AS total_orders FROM orders WHERE status = 'completed' AND order_date LIKE '2024%'",
+    },
+    2: {
+        "description": "Find the top 5 customers by total revenue (sum of total_amount for completed orders only). Return columns: first_name, last_name, total_revenue. Order by total_revenue descending.",
+        "difficulty": "medium",
+        "hint": "JOIN orders with customers, GROUP BY customer, filter completed, ORDER and LIMIT",
+        "answer_query": "SELECT c.first_name, c.last_name, ROUND(SUM(o.total_amount), 2) AS total_revenue FROM orders o JOIN customers c ON o.customer_id = c.customer_id WHERE o.status = 'completed' GROUP BY o.customer_id ORDER BY total_revenue DESC LIMIT 5",
+    },
+    3: {
+        "description": "For each product category, calculate the total revenue (completed orders only) and rank categories by revenue using a window function. Return columns: category, total_revenue, revenue_rank. Order by revenue_rank ascending.",
+        "difficulty": "hard",
+        "hint": "Use SUM with GROUP BY inside a CTE, then apply RANK() OVER (ORDER BY ...) on the result",
+        "answer_query": "WITH category_revenue AS ( SELECT p.category, SUM(o.total_amount) AS total_revenue FROM orders o JOIN products p ON o.product_id = p.product_id WHERE o.status = 'completed' GROUP BY p.category ) SELECT category, total_revenue, RANK() OVER (ORDER BY total_revenue DESC) AS revenue_rank FROM category_revenue ORDER BY revenue_rank ASC",
+    },
+    4: {
+        "description": "Find the average price of products in each category, but only for categories that have more than 2 products. Return columns: category, avg_price.",
+        "difficulty": "medium",
+        "hint": "Use GROUP BY with HAVING COUNT(...) > 2.",
+        "answer_query": "SELECT category, ROUND(AVG(price), 2) AS avg_price FROM products GROUP BY category HAVING COUNT(product_id) > 2",
+    },
+    5: {
+        "description": "Identify customers who have ordered products from both the 'Electronics' and 'Clothing' categories. Return columns: customer_id, first_name.",
+        "difficulty": "hard",
+        "hint": "Use INTERSECT on two queries, or GROUP BY customer HAVING COUNT(DISTINCT category) = 2.",
+        "answer_query": "SELECT DISTINCT c.customer_id, c.first_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN products p ON o.product_id = p.product_id WHERE p.category = 'Electronics' INTERSECT SELECT DISTINCT c.customer_id, c.first_name FROM customers c JOIN orders o ON c.customer_id = o.customer_id JOIN products p ON o.product_id = p.product_id WHERE p.category = 'Clothing'",
+    },
+}
+
+# ── Sessions ──────────────────────────────────────────────────────────────────
+
+sessions = {}
+
+def get_session(sid: str):
+    if sid not in sessions:
+        sessions[sid] = {
+            "task_id": None, "task": None, 
+            "expected_rows": None, "expected_columns": None,
+            "attempts": 0, "best_reward": 0.0, "history": []
+        }
+    return sessions[sid]
+
+# ── Database Helpers ──────────────────────────────────────────────────────────
+
+def progress_handler():
+    raise sqlite3.OperationalError("Query execution aborted: Timed out or exceeded instruction limits. Hint: Too complex CROSS JOIN?")
+
+def get_connection():
+    if not os.path.exists(DB_PATH):
+        raise HTTPException(status_code=500, detail=f"Database not found at {DB_PATH}.")
+    conn = sqlite3.connect(DB_PATH)
+    conn.row_factory = sqlite3.Row
+    # Security feature: Prevent DOS
+    conn.set_progress_handler(progress_handler, 500000)
+    return conn
+
+def get_schema_info() -> str:
+    conn = get_connection()
+    cur = conn.cursor()
+    schema_parts = []
+    cur.execute("SELECT name FROM sqlite_master WHERE type='table'")
+    tables = [r["name"] for r in cur.fetchall()]
+    for table in tables:
+        cur.execute(f"PRAGMA table_info({table})")
+        cols = cur.fetchall()
+        col_defs = ", ".join(f"{c['name']} {c['type']}" for c in cols)
+        cur.execute(f"SELECT COUNT(*) AS n FROM {table}")
+        count = cur.fetchone()["n"]
+        schema_parts.append(f"  {table} ({col_defs})  -- {count} rows")
+    conn.close()
+    return "Tables:\\n" + "\\n".join(schema_parts)
+
+def run_query(sql: str) -> tuple[list[dict], list[str]]:
+    conn = get_connection()
+    cur = conn.cursor()
+    cur.execute(sql)
+    columns = [d[0] for d in cur.description] if cur.description else []
+    rows = [dict(zip(columns, row)) for row in cur.fetchall()]
+    conn.close()
+    return rows, columns
+
+def compute_expected(sid: str):
+    session = get_session(sid)
+    task = session["task"]
+    rows, columns = run_query(task["answer_query"])
+    session["expected_rows"] = rows
+    session["expected_columns"] = columns
+
+# ── Reward Function ───────────────────────────────────────────────────────────
+
+def compute_reward(sid: str, agent_rows: list[dict], agent_cols: list[str]) -> tuple[float, dict]:
+    session = get_session(sid)
+    expected_rows = session["expected_rows"]
+    expected_cols = session["expected_columns"]
+    details = {}
+
+    agent_cols_lower   = [c.lower() for c in agent_cols]
+    expected_cols_lower = [c.lower() for c in expected_cols]
+    col_matches = sum(1 for c in expected_cols_lower if c in agent_cols_lower)
+    col_score = (col_matches / len(expected_cols_lower)) * 0.30 if expected_cols_lower else 0.0
+    details["column_score"] = round(col_score, 3)
+
+    expected_count = len(expected_rows)
+    agent_count    = len(agent_rows)
+    if expected_count == 0:
+        row_score = 0.30 if agent_count == 0 else 0.0
+    else:
+        row_ratio = min(agent_count, expected_count) / max(agent_count, expected_count)
+        row_score = row_ratio * 0.30
+    details["row_score"] = round(row_score, 3)
+
+    if not expected_rows or not agent_rows:
+        value_score = 0.0
+    else:
+        def normalize(v):
+            if v is None: return ""
+            try: return str(round(float(v), 1))
+            except: return str(v).strip().lower()
+
+        matched_cells = 0
+        total_cells   = len(expected_rows) * len(expected_cols_lower)
+
+        for exp_row, agt_row in zip(expected_rows, agent_rows):
+            for col in expected_cols_lower:
+                exp_val = normalize(exp_row.get(col) or exp_row.get(col.upper()))
+                agt_val = normalize(agt_row.get(col) or agt_row.get(col.upper()))
+                if not agt_val:
+                    exp_idx = expected_cols_lower.index(col)
+                    if exp_idx < len(agent_cols):
+                        pos_col = agent_cols[exp_idx]
+                        agt_val = normalize(agt_row.get(pos_col))
+                if exp_val == agt_val:
+                    matched_cells += 1
+
+        value_score = (matched_cells / total_cells) * 0.40 if total_cells > 0 else 0.0
+    details["value_score"] = round(value_score, 3)
+
+    total = round(col_score + row_score + value_score, 3)
+    details["total_reward"] = total
+    return total, details
+
+# ── Endpoints ─────────────────────────────────────────────────────────────────
+
+@api_app.post("/reset", response_model=ResetResponse)
+def reset(req: ResetRequest):
+    if req.task_id not in TASKS: raise HTTPException(status_code=400, detail="Invalid task_id")
+    session = get_session(req.session_id)
+    session["task_id"]    = req.task_id
+    session["task"]       = TASKS[req.task_id]
+    session["attempts"]   = 0
+    session["best_reward"] = 0.0
+    session["history"]    = []
+    compute_expected(req.session_id)
+    observation = {
+        "task_id":          req.task_id,
+        "difficulty":       session["task"]["difficulty"],
+        "task_description": session["task"]["description"],
+        "schema":           get_schema_info(),
+        "hint":             session["task"]["hint"],
+    }
+    return ResetResponse(observation=observation, info={"message": f"Task {req.task_id} loaded."})
+
+@api_app.post("/step", response_model=StepResponse)
+def step(req: StepRequest):
+    session = get_session(req.session_id)
+    if session["task_id"] is None: raise HTTPException(status_code=400, detail="Call /reset first.")
+    session["attempts"] += 1
+    sql = req.action.strip()
+
+    if not re.match(r"^\\s*(SELECT|WITH)\\b", sql, re.IGNORECASE):
+        return StepResponse(
+            observation={"error": "Only SELECT or WITH allowed."}, reward=0.0, done=False,
+            info={"attempt": session["attempts"], "message": "Rejected"}
+        )
+
+    try:
+        agent_rows, agent_cols = run_query(sql)
+    except Exception as e:
+        session["history"].append({"attempt": session["attempts"], "sql": sql, "reward": 0.0, "error": str(e)})
+        return StepResponse(
+            observation={"error": str(e), "sql_submitted": sql}, reward=0.0, done=False,
+            info={"attempt": session["attempts"], "message": "SQL Error"}
+        )
+
+    reward, details = compute_reward(req.session_id, agent_rows, agent_cols)
+    session["best_reward"] = max(session["best_reward"], reward)
+    done = reward >= 1.0
+    session["history"].append({"attempt": session["attempts"], "sql": sql, "reward": reward, "details": details})
+
+    observation = {
+        "task_id": session["task_id"], "task_description": session["task"]["description"],
+        "sql_submitted": sql, "result_preview": agent_rows[:5], "result_row_count": len(agent_rows),
+        "reward_breakdown": details,
+    }
+    return StepResponse(
+        observation=observation, reward=reward, done=done,
+        info={"attempt": session["attempts"], "best_reward": session["best_reward"]}
+    )
+
+@api_app.get("/state", response_model=StateResponse)
+def state(req: StateRequest):
+    session = get_session(req.session_id)
+    return StateResponse(
+        session_id=req.session_id,
+        task_id=session["task_id"],
+        task_description=session["task"]["description"] if session["task"] else None,
+        schema_info=get_schema_info(),
+        attempts=session["attempts"],
+        best_reward=session["best_reward"],
+        history=session["history"],
+    )
+
+@api_app.get("/")
+def root():
+    return {"message": "API running at /api. Try the UI at root (handled by wrapper)!"}
+
+# ── Server Setup ──────────────────────────────────────────────────────────────
+
+demo = build_ui()
+app = gr.mount_gradio_app(api_app, demo, path="/")
+"""
+with open("main.py", "w", encoding="utf-8") as f:
+    f.write(main_py_code.strip() + "\n")
+
+# 4. Modify inference.py to use session IDs
+with open("inference.py", "r", encoding="utf-8") as f:
+    inf = f.read()
+
+inf = inf.replace('ENV_BASE_URL = "http://127.0.0.1:7860"', 'ENV_BASE_URL = "http://127.0.0.1:7860/api"')
+inf = inf.replace("TASK_IDS     = [1, 2, 3]", "TASK_IDS     = [1, 2, 3, 4, 5]")
+inf = inf.replace('def env_reset(task_id: int) -> dict:', 'def env_reset(task_id: int, session_id: str) -> dict:')
+inf = inf.replace('json={"task_id": task_id}', 'json={"task_id": task_id, "session_id": session_id}')
+
+inf = inf.replace('def env_step(sql: str) -> dict:', 'def env_step(sql: str, session_id: str) -> dict:')
+inf = inf.replace('json={"action": sql}', 'json={"action": sql, "session_id": session_id}')
+
+inf = inf.replace("reset_resp = env_reset(task_id)", 'session_id = f"baseline_{task_id}"\n    reset_resp = env_reset(task_id, session_id)')
+inf = inf.replace("step_resp = env_step(sql)", 'step_resp = env_step(sql, session_id)')
+
+inf = inf.replace('r = requests.get(f"{ENV_BASE_URL}/state")', 'r = requests.get(f"{ENV_BASE_URL}/state", json={"session_id": "baseline_1"})')
+
+# health check fix for inference wait_for_server
+inf = re.sub(r'requests\.get\(f"\{ENV_BASE_URL\}/health", timeout=3\)', 'requests.get(f"{ENV_BASE_URL}/", timeout=3)', inf)
+
+with open("inference.py", "w", encoding="utf-8") as f:
+    f.write(inf)
+
+print("Done generating make_awesome.")

requirements.txt

@@ -3,4 +3,7 @@ uvicorn==0.30.6
 pydantic==2.9.2
 requests==2.32.3
 openai==1.51.0
-pyyaml==6.0.2
+python-dotenv==1.0.1
+pyyaml==6.0.2
+gradio==4.44.0
+pandas

ui.py

@@ -0,0 +1,86 @@
+import gradio as gr
+import requests
+import pandas as pd
+import json
+import uuid
+
+ENV_BASE_URL = "http://127.0.0.1:7860"
+
+def new_session():
+    return str(uuid.uuid4())
+
+def load_task(task_id, sid):
+    try:
+        task_num = int(task_id.split()[1])
+        r = requests.post(f"{ENV_BASE_URL}/api/reset", json={"task_id": task_num, "session_id": sid})
+        if r.status_code != 200:
+            return f"Error: {r.text}", "", "", pd.DataFrame(), f"Error loading task {task_num}"
+            
+        data = r.json()
+        obs = data["observation"]
+        return obs["task_description"], obs["schema"], obs["hint"], pd.DataFrame(), "Task loaded. Write SQL below."
+    except Exception as e:
+        return str(e), "", "", pd.DataFrame(), "Error loading task"
+
+def run_sql(sql, sid):
+    if not sql.strip():
+        return pd.DataFrame(), "Please enter a SQL query.", "Error"
+    try:
+        r = requests.post(f"{ENV_BASE_URL}/api/step", json={"action": sql, "session_id": sid})
+        data = r.json()
+        obs = data.get("observation", {})
+        reward = data.get("reward", 0.0)
+        done = data.get("done", False)
+        
+        df = pd.DataFrame(obs.get("result_preview", []))
+        breakdown = obs.get("reward_breakdown", {})
+        
+        feedback = f"🎯 Reward: {reward:.2f} / 1.0\n"
+        if breakdown:
+            feedback += f"Columns: {breakdown.get('column_score',0):.2f}, Rows: {breakdown.get('row_score',0):.2f}, Values: {breakdown.get('value_score',0):.2f}"
+        if "error" in obs:
+            feedback += f"\n\n⚠️ Error: {obs['error']}"
+        
+        return df, feedback, "✅ SOLVED!" if done else "Keep trying!"
+    except Exception as e:
+         return pd.DataFrame(), str(e), "Error"
+
+def build_ui():
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
+        gr.Markdown("# 📊 SQL Analyst OpenEnv")
+        gr.Markdown("An interactive environment to test SQL generation. Load a task, read the schema, and write a query answering the business question.")
+        
+        sid = gr.State(new_session)
+        
+        with gr.Group():
+            gr.Markdown("### Step 1: Select a Task")
+            with gr.Row():
+                task_dropdown = gr.Dropdown(choices=["Task 1 (Easy)", "Task 2 (Medium)", "Task 3 (Hard)", "Task 4 (Medium)", "Task 5 (Hard)"], value="Task 1 (Easy)", label="Available Tasks", show_label=False)
+                btn_load = gr.Button("🔄 Load Task", variant="primary")
+        
+        with gr.Row():
+            desc = gr.Textbox(label="🎯 Business Question", interactive=False, lines=2)
+            
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Step 2: Understand the Data")
+                schema = gr.Code(label="Database Schema", language="sql", interactive=False)
+                with gr.Accordion("💡 Need a hint?", open=False):
+                    hint = gr.Textbox(show_label=False, interactive=False)
+            
+            with gr.Column():
+                gr.Markdown("### Step 3: Write & Execute SQL")
+                sql_input = gr.Code(label="SQL Editor", language="sql", lines=12)
+                btn_run = gr.Button("🚀 Execute Query", variant="primary")
+                
+                status_out = gr.Markdown("Waiting for query...")
+                feedback_out = gr.Textbox(label="Evaluation Score & Feedback", interactive=False, lines=3)
+                
+        with gr.Group():
+            gr.Markdown("### Step 4: Review Results")
+            grid_out = gr.Dataframe(label="Result Preview (First 5 Rows)", interactive=False)
+
+        btn_load.click(load_task, inputs=[task_dropdown, sid], outputs=[desc, schema, hint, grid_out, feedback_out])
+        btn_run.click(run_sql, inputs=[sql_input, sid], outputs=[grid_out, feedback_out, status_out])
+        
+    return demo