Massive Rewrite: Rebuilt Phase 2 standalone validation constraints perfectly mapped to 0.05-0.95

inference.py

@@ -1,260 +1,164 @@
-"""
-SQL Data Analyst OpenEnv — baseline inference (hackathon format).
-
-Environment variables (set before running):
-  HF_TOKEN         Primary API key (Hugging Face / OpenAI-compatible).
-  API_BASE_URL     LLM base URL (e.g. https://api.openai.com/v1 or HF router).
-  MODEL_NAME       Model id for chat completions.
-  OPENAI_API_KEY   Optional fallback if HF_TOKEN is unset.
-  API_KEY          Optional second fallback.
-
-Optional:
-  SQL_AGENT_TASK      Logged as task= in [START] (default: sql_analyst_episode).
-  SQL_AGENT_BENCHMARK Logged as env= in [START] (default: sql_agent_openenv).
-  MAX_STEPS           Max env.step calls (default: 24).
-  OPENAI_SEED         Passed to OpenAI only when base URL looks like OpenAI.
-  ENV_CONTAINER_START true to use SqlEnvClient.from_docker_image (default: false).
-  IMAGE_NAME / LOCAL_IMAGE_NAME / ENV_IMAGE_NAME  Docker image tag when using container.
-"""
-
-from __future__ import annotations
-
 import os
 import re
-import sys
-import textwrap
-from typing import Any, Dict, List, Optional
-
-from dotenv import load_dotenv
+import json
+import requests
 from openai import OpenAI
 
-load_dotenv()
-
-from client import SqlEnvClient
-from models import SqlAction
-
-# --- Mandatory hackathon configuration ---
-API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
-MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
-API_KEY = (
-    os.getenv("HF_TOKEN")
-    or os.getenv("OPENAI_API_KEY")
-    or os.getenv("API_KEY")
-)
-
-TASK_NAME = os.getenv("SQL_AGENT_TASK", "sql_analyst_episode")
-BENCHMARK = os.getenv("SQL_AGENT_BENCHMARK", "sql_agent_openenv")
-
-MAX_STEPS = int(os.getenv("MAX_STEPS", "24"))
-OPENAI_SEED = int(os.getenv("OPENAI_SEED", "42"))
-
-ENV_CONTAINER_START = os.getenv("ENV_CONTAINER_START", "false").lower() == "true"
-ENV_IMAGE_NAME = (
-    os.getenv("LOCAL_IMAGE_NAME")
-    or os.getenv("IMAGE_NAME")
-    or os.getenv("ENV_IMAGE_NAME", "sql-agent-env:latest")
-)
-
-# Grader task ids (must match core/tasks.py) for normalized [END] score in [0, 1]
-TASK_IDS = [
-    "employee_payroll_overview",
-    "department_budget_summary",
-    "senior_engineering_comp_review",
-]
-
-SYSTEM_PROMPT = textwrap.dedent(
-    """
-    You are an expert Data Analyst and SQL Agent.
-    You will be provided with a SQL Schema and an instruction.
-    Reply with exactly one SQL query string to execute.
-    Do NOT use markdown fences or explanations — only the raw SQL text.
-    """
-).strip()
-
-
-def log_start(task: str, env: str, model: str) -> None:
-    print(f"[START] task={task} env={env} model={model}")
-
-def log_step(
-    step: int,
-    action: str,
-    reward: float,
-    done: bool,
-    error: Optional[str],
-) -> None:
-    err_one = sanitize_one_line(error) if error else "null"
-    act_one = sanitize_one_line(action)
-    done_val = "true" if done else "false"
-    print(f"[STEP]  step={step} action={act_one} reward={reward:.2f} done={done_val} error={err_one}")
-
-def log_end(success: bool, steps: int, score: float) -> None:
-    succ_val = "true" if success else "false"
-    print(f"[END]   success={succ_val} steps={steps} rewards={score:.2f}")
-
-
-def sanitize_one_line(s: str) -> str:
-    if not s:
-        return ""
-    return " ".join(s.split())
-
-
-def build_user_prompt(step: int, observation: Any, history: List[str]) -> str:
-    schema = observation.schema_info
-    instruction = observation.current_task_instruction
-    exec_result = observation.execution_result or "None"
-    error = observation.execution_error or "None"
-    history_block = "\n".join(history[-6:]) if history else "None"
-    return textwrap.dedent(
-        f"""
-        Step: {step}
-        Database Schema:
-        {schema}
-
-        Task Instruction:
-        {instruction}
-
-        Previous Execution Result: {exec_result}
-        Previous Error: {error}
-
-        Recent history:
-        {history_block}
-
-        Write the precise SQL query string to accomplish the task instruction. Return nothing else.
-        """
-    ).strip()
-
-
-def parse_model_action(response_text: str) -> str:
-    query = (response_text or "").strip()
-    if query.startswith("```sql"):
-        query = query[6:]
-    if query.startswith("```"):
-        query = query[3:]
-    if query.endswith("```"):
-        query = query[:-3]
-    return query.strip()
-
-
-def _make_direct_client():
-    from server.environment import SqlEnvironment
-
-    base_env = SqlEnvironment()
-
-    class DirectClient:
-        def __init__(self, target):
-            self.target = target
-
-        def reset(self):
-            obs = self.target.reset()
-            return type(
-                "StepResult",
-                (),
-                {"observation": obs, "reward": 0.0, "done": False},
-            )()
-
-        def step(self, action):
-            obs = self.target.step(action)
-            return type(
-                "StepResult",
-                (),
-                {
-                    "observation": obs,
-                    "reward": obs.reward if obs.reward is not None else 0.0,
-                    "done": obs.done,
-                },
-            )()
-
-        def close(self):
-            pass
-
-    return DirectClient(base_env)
-
-
-def main() -> None:
-    if not API_KEY:
-        sys.exit(1)
-
-    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-
-    if ENV_CONTAINER_START:
-        env = SqlEnvClient.from_docker_image(ENV_IMAGE_NAME).sync()
-    else:
-        env = _make_direct_client()
-
-    history: List[str] = []
+API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "")
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+ENV_URL = os.getenv("ENV_URL", "http://localhost:8000") # defaults to common uvicorn port but user specified 7860
+MAX_STEPS = 6
+TEMPERATURE = 0.1
+MAX_TOKENS = 500
+
+client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+
+SYSTEM_PROMPT = """You are an expert SQL developer.
+For write_query tasks: write a correct SQL SELECT query.
+For fix_query tasks: fix the broken SQL provided.
+For optimize_query tasks: rewrite slow SQL using CTEs or JOINs instead of subqueries.
+
+ALWAYS respond with ONLY a JSON object like this:
+{"action_type": "write_query", "sql": "SELECT ...", "explanation": "reason"}
+
+Rules: Only SELECT allowed. No DROP DELETE INSERT UPDATE CREATE ALTER."""
+
+def build_prompt(obs: dict) -> str:
+    obs_data = obs.get("observation", obs)
+    task_desc = obs_data.get("task_description", "")
+    schema_info = obs_data.get("schema_info", "")
+    sample_data = obs_data.get("sample_data", "")
+    hints = obs_data.get("metadata", [])
+    last_sql = obs_data.get("last_sql", "")
+    last_result = obs_data.get("last_result", "")
+    last_error = obs_data.get("last_error", "")
+    step_count = obs_data.get("step_count", 0)
+    feedback = obs_data.get("feedback", "")
     
-    try:
-        result = env.reset()
-        observation = result.observation
+    prompt = f"Task Description:\n{task_desc}\n\nSchema Info:\n{schema_info}\n\nSample Data:\n{sample_data}\n"
+    if hints:
+        prompt += f"\nHints: {', '.join(hints)}\n"
+    if last_sql:
+        prompt += f"\nLast SQL Submitted: {last_sql}\n"
+    if last_result:
+        prompt += f"Result of Last SQL: {last_result}\n"
+    if last_error:
+        prompt += f"Error Message: {last_error}\n"
+    if step_count > 0 and feedback:
+        prompt += f"Feedback from Grader: {feedback}\n"
         
-        current_task_id = observation.task_id
-        task_step = 1
-        log_start(task=current_task_id, env=BENCHMARK, model=MODEL_NAME)
-
-        for step in range(1, MAX_STEPS + 1):
-            if result.done:
-                break
-
-            user_prompt = build_user_prompt(step, observation, history)
-            messages = [
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
-            ]
-
-            create_kwargs: Dict[str, Any] = {
-                "model": MODEL_NAME,
-                "messages": messages,
-                "temperature": 0.0,
-                "stream": False,
-            }
-            if re.search(r"openai\.com", API_BASE_URL, re.I):
-                create_kwargs["seed"] = OPENAI_SEED
+    prompt += "\nRespond with JSON action only."
+    return prompt
 
+def parse_action(response_text: str, task_type: str) -> dict:
+    try:
+        return json.loads(response_text)
+    except json.JSONDecodeError:
+        json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+        if json_match:
             try:
-                completion = client.chat.completions.create(**create_kwargs)
-                response_text = completion.choices[0].message.content or "SELECT 1"
-            except Exception:
-                response_text = "SELECT 1"
+                return json.loads(json_match.group(0))
+            except json.JSONDecodeError:
+                pass
+        
+        sql_match = re.search(r'(?i)SELECT\s+.*', response_text, re.DOTALL)
+        sql = sql_match.group(0).strip() if sql_match else "SELECT 1"
+        if sql.endswith('```'): sql = sql[:-3].strip()
+        
+        return {"action_type": task_type, "sql": sql, "explanation": "Fallback parsed"}
 
-            action_str = parse_model_action(response_text)
-            result = env.step(SqlAction(query=action_str))
-            observation = result.observation
+def run_episode(task_id: str) -> dict:
+    res = requests.post(f"{ENV_URL}/reset", json={"task_id": task_id})
+    if res.status_code != 200:
+        return {"task_id": task_id, "best_reward": 0.05}
+    obs_obj = res.json()
+    
+    best_reward = 0.05
+    for step in range(MAX_STEPS):
+        obs = obs_obj.get("observation", obs_obj)
+        if obs.get("done", False):
+            break
             
-            # reward is obtained from observation task_score or directly from step
-            task_score = float(observation.task_score if hasattr(observation, 'task_score') and observation.task_score is not None else 0.0)
-            if hasattr(observation, 'reward') and observation.reward is not None:
-                task_score = float(observation.reward)
-                
-            err_raw = observation.execution_error
-            next_task_id = getattr(observation, "task_id", None)
+        prompt = build_prompt(obs_obj)
+        
+        try:
+            completion = client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=[
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": prompt}
+                ],
+                temperature=TEMPERATURE,
+                max_tokens=MAX_TOKENS
+            )
+            response_text = completion.choices[0].message.content or "{}"
+        except Exception:
+            response_text = "{}"
             
-            # task transitions or episode done signals end of current task
-            task_done = (next_task_id != current_task_id) or result.done
+        action = parse_action(response_text, obs.get("task_type", "write_query"))
+        
+        res = requests.post(f"{ENV_URL}/step", json=action)
+        if res.status_code != 200:
+            break
+        obs_obj = res.json()
+        reward = obs_obj.get("reward", 0.05)
+        
+        best_reward = max(best_reward, reward)
+        print(f"Step {step+1}: SQL={action.get('sql', '')[:60]}... Reward={reward:.3f} Feedback={obs_obj.get('observation', {}).get('feedback', '')}")
+        
+        if obs_obj.get("done", obs_obj.get("observation", {}).get("done", False)):
+            break
             
-            log_step(step=task_step, action=action_str, reward=task_score, done=task_done, error=err_raw)
-
-            history.append(f"Q: {action_str[:200]} -> R: {task_score:+.2f}")
-
-            if task_done:
-                success = (task_score >= 0.95)
-                log_end(success=success, steps=task_step, score=task_score)
-                if result.done or not next_task_id:
-                    break
-                current_task_id = next_task_id
-                task_step = 1
-                log_start(task=current_task_id, env=BENCHMARK, model=MODEL_NAME)
-                history.clear()
-            else:
-                task_step += 1
+    return {"task_id": task_id, "best_reward": round(best_reward, 3)}
 
+def main():
+    print(f"--- SQL Agent Inference ---")
+    print(f"Model: {MODEL_NAME}")
+    print(f"Env URL: {ENV_URL}")
+    
+    health = {"status": "unreachable"}
+    try:
+        health = requests.get(f"{ENV_URL}/health").json()
     except Exception:
         pass
-    finally:
+    print(f"Health: {health}")
+    
+    task_ids = ["easy_01", "easy_02", "medium_01", "medium_02", "hard_01", "hard_02"]
+    results = []
+    
+    for tid in task_ids:
+        print(f"\n=== Task: {tid} ===")
         try:
-            env.close()
-        except Exception:
-            pass
-
+            res = run_episode(tid)
+        except Exception as e:
+            res = {"task_id": tid, "best_reward": 0.05}
+            print(f"Error running task {tid}: {e}")
+        results.append(res)
+        
+    print("\n=== FINAL SCORES ===")
+    easy_scores = []
+    medium_scores = []
+    hard_scores = []
+    
+    for r in results:
+        t = r["task_id"]
+        v = r["best_reward"]
+        print(f"{t}: {v}")
+        if t.startswith("easy"): easy_scores.append(v)
+        elif t.startswith("medium"): medium_scores.append(v)
+        elif t.startswith("hard"): hard_scores.append(v)
+        
+    easy_avg = sum(easy_scores)/len(easy_scores) if easy_scores else 0.0
+    medium_avg = sum(medium_scores)/len(medium_scores) if medium_scores else 0.0
+    hard_avg = sum(hard_scores)/len(hard_scores) if hard_scores else 0.0
+    overall = (easy_avg + medium_avg + hard_avg) / 3.0
+    
+    print("\n--- AVERAGES ---")
+    print(f"Easy Average:   {easy_avg:.3f}")
+    print(f"Medium Average: {medium_avg:.3f}")
+    print(f"Hard Average:   {hard_avg:.3f}")
+    print(f"Overall Score:  {overall:.3f}")
 
 if __name__ == "__main__":
     main()

openenv.yaml

@@ -1,8 +1,9 @@
 name: sql-debugger-env
 version: "1.0.0"
 description: >
-  RL environment for training AI agents to write, debug, and optimize SQL queries
-  against a real SQLite database with employees, departments, and projects.
+  An RL environment for training AI agents to write, fix, and optimize
+  SQL queries against a real SQLite database with employees, departments,
+  projects and project assignments tables.
 
 tasks:
   - id: easy_01
@@ -10,64 +11,71 @@ tasks:
     difficulty: easy
     type: write_query
     grader: true
-    description: "Find Engineering employees with salary above 90000"
-  
+    description: "Find Engineering employees with salary above 90000, return name and salary ordered by salary DESC"
+
   - id: easy_02
-    name: Count per department
+    name: Count by department
     difficulty: easy
     type: write_query
     grader: true
-    description: "Count employees per department, return department and count ordered by count DESC."
+    description: "Count employees per department, return department and count ordered by count DESC"
 
   - id: medium_01
     name: Fix broken JOIN
     difficulty: medium
     type: fix_query
     grader: true
-    description: "Fix a broken JOIN query (missing ON keyword, wrong table alias in WHERE). Return name and department name."
+    description: "Fix the broken JOIN query missing ON keyword and using wrong table alias in WHERE"
 
   - id: medium_02
-    name: Fix GROUP BY
+    name: Fix wrong GROUP BY
     difficulty: medium
     type: fix_query
     grader: true
-    description: "Fix wrong GROUP BY (should GROUP BY department not id) finding max salary per department."
+    description: "Fix the GROUP BY clause that incorrectly groups by id instead of department"
 
   - id: hard_01
     name: Optimize correlated subquery
     difficulty: hard
     type: optimize_query
     grader: true
-    description: "Replace correlated subquery with window function/CTE to find top earner per department with total hours."
+    description: "Replace correlated subqueries with window functions or CTEs to find top earner per department"
 
   - id: hard_02
-    name: Eliminate N+1 subqueries
+    name: Eliminate N+1 problem
     difficulty: hard
     type: optimize_query
     grader: true
-    description: "Eliminate N+1 subqueries using LEFT JOIN + GROUP BY to compute total hours per project."
+    description: "Rewrite N+1 correlated subquery using LEFT JOIN and GROUP BY"
 
 action_space:
   type: object
-  fields:
-    action_type: string
-    sql: string
-    explanation: string
+  properties:
+    action_type:
+      type: string
+      enum: [write_query, fix_query, optimize_query]
+    sql:
+      type: string
+      description: SQL SELECT statement
+    explanation:
+      type: string
+      description: Optional reasoning
 
 observation_space:
   type: object
-  fields:
-    task_id: string
-    task_type: string
-    task_description: string
-    schema_info: string
-    last_sql: string
-    last_result: string
-    last_error: string
-    step_count: integer
-    done: boolean
-    reward: float
-    feedback: string
+  properties:
+    task_id: {type: string}
+    task_type: {type: string}
+    task_description: {type: string}
+    schema_info: {type: string}
+    sample_data: {type: string}
+    last_sql: {type: string}
+    last_result: {type: string}
+    last_error: {type: string}
+    step_count: {type: integer}
+    done: {type: boolean}
+    reward: {type: number, minimum: 0.05, maximum: 0.95}
+    feedback: {type: string}
 
-reward_range: [0.0, 1.0]
+reward_range: [0.05, 0.95]
 max_steps_per_episode: 8

server/app.py

@@ -1,28 +1,34 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import Optional, Dict, Any
 
-from server.environment import SqlEnvironment, TASKS
+from server.environment import SQLEnvironment
 
-app = FastAPI(title="SQL Agent Environment API")
-
-env = SqlEnvironment()
+app = FastAPI(title="SQL Debugger OpenEnv")
+env = SQLEnvironment()
 
 class ResetRequest(BaseModel):
     task_id: Optional[str] = None
     difficulty: Optional[str] = None
 
-class ActionRequest(BaseModel):
-    sql: str
-    action_type: Optional[str] = None
+class StepRequest(BaseModel):
+    action_type: str = "write_query"
+    sql: str = ""
+    explanation: Optional[str] = None
+    metadata: dict = {}
 
 class GradeRequest(BaseModel):
     task_id: str
-    action: ActionRequest
+    action: dict
 
 @app.get("/")
 def read_root():
-    return {"info": "SQL Agent Environment API"}
+    return {
+        "info": "SQL Agent Environment API",
+        "version": "1.0.0",
+        "tasks": 6,
+        "status": "running"
+    }
 
 @app.get("/health")
 def read_health():
@@ -32,10 +38,6 @@ def read_health():
 def read_state():
     return env.state
 
-@app.get("/tasks")
-def list_tasks():
-    return {"tasks": TASKS}
-
 @app.post("/reset")
 def do_reset(req: Optional[ResetRequest] = None):
     if req:
@@ -45,19 +47,78 @@ def do_reset(req: Optional[ResetRequest] = None):
     return obs
 
 @app.post("/step")
-def do_step(action: ActionRequest):
-    obs = env.step({"sql": action.sql, "action_type": action.action_type})
+def do_step(req: StepRequest):
+    obs = env.step(req.model_dump() if hasattr(req, "model_dump") else req.dict())
     return obs
 
+@app.get("/tasks")
+def list_tasks():
+    return {
+        "tasks": [
+            {
+                "id": "easy_01",
+                "name": "Filter high earners",
+                "difficulty": "easy",
+                "type": "write_query",
+                "grader": True,
+                "description": "Find all employees in Engineering with salary above 90000"
+            },
+            {
+                "id": "easy_02",
+                "name": "Count by department",
+                "difficulty": "easy",
+                "type": "write_query",
+                "grader": True,
+                "description": "Count employees per department ordered by count"
+            },
+            {
+                "id": "medium_01",
+                "name": "Fix broken JOIN",
+                "difficulty": "medium",
+                "type": "fix_query",
+                "grader": True,
+                "description": "Fix the broken JOIN query for active project assignments"
+            },
+            {
+                "id": "medium_02",
+                "name": "Fix wrong GROUP BY",
+                "difficulty": "medium",
+                "type": "fix_query",
+                "grader": True,
+                "description": "Fix the GROUP BY clause to aggregate by department"
+            },
+            {
+                "id": "hard_01",
+                "name": "Optimize correlated subquery",
+                "difficulty": "hard",
+                "type": "optimize_query",
+                "grader": True,
+                "description": "Replace correlated subqueries with window functions"
+            },
+            {
+                "id": "hard_02",
+                "name": "Eliminate N+1 problem",
+                "difficulty": "hard",
+                "type": "optimize_query",
+                "grader": True,
+                "description": "Rewrite N+1 query using LEFT JOIN and GROUP BY"
+            }
+        ]
+    }
+
 @app.post("/grade")
 def do_grade(req: GradeRequest):
     env.reset(task_id=req.task_id)
-    obs = env.step({"sql": req.action.sql, "action_type": req.action.action_type})
-    score = max(0.05, min(0.95, float(obs["reward"])))
+    step_result = env.step(req.action)
+    
+    raw_score = step_result.get("reward", 0.05)
+    score = max(0.05, min(0.95, float(raw_score)))
+    score = round(score, 4)
     
     return {
         "task_id": req.task_id,
         "score": score,
-        "feedback": obs["feedback"],
-        "done": obs["done"]
+        "feedback": step_result["observation"]["feedback"],
+        "done": step_result["done"],
+        "info": step_result.get("info", {})
     }

server/environment.py

@@ -4,257 +4,325 @@ from typing import Dict, Any, List, Optional
 import json
 
 SCHEMA_SQL = """
-CREATE TABLE employees (id INTEGER PRIMARY KEY, name TEXT, department TEXT, salary INTEGER, hire_date TEXT, manager_id INTEGER);
-CREATE TABLE departments (id INTEGER PRIMARY KEY, name TEXT, budget INTEGER, location TEXT);
+CREATE TABLE employees (id INTEGER PRIMARY KEY, name TEXT, department TEXT, salary REAL, hire_date TEXT, manager_id INTEGER);
+CREATE TABLE departments (id INTEGER PRIMARY KEY, name TEXT, budget REAL, location TEXT);
 CREATE TABLE projects (id INTEGER PRIMARY KEY, name TEXT, department_id INTEGER, start_date TEXT, status TEXT);
-CREATE TABLE project_assignments (employee_id INTEGER, project_id INTEGER, hours_worked INTEGER);
+CREATE TABLE project_assignments (employee_id INTEGER, project_id INTEGER, hours_worked REAL, PRIMARY KEY(employee_id, project_id));
 """
 
 SEED_DATA_SQL = """
-INSERT INTO employees VALUES (1, 'Alice', 'Engineering', 95000, '2020-01-15', 3);
-INSERT INTO employees VALUES (2, 'Bob', 'Engineering', 85000, '2021-02-20', 3);
-INSERT INTO employees VALUES (3, 'Charlie', 'Engineering', 120000, '2019-01-10', NULL);
-INSERT INTO employees VALUES (4, 'David', 'Sales', 70000, '2022-03-05', 5);
-INSERT INTO employees VALUES (5, 'Eve', 'Sales', 92000, '2018-11-20', NULL);
-INSERT INTO employees VALUES (6, 'Frank', 'HR', 65000, '2023-05-12', NULL);
-
-INSERT INTO departments VALUES (1, 'Engineering', 500000, 'Building A');
-INSERT INTO departments VALUES (2, 'Sales', 300000, 'Building B');
-INSERT INTO departments VALUES (3, 'HR', 150000, 'Building C');
-
-INSERT INTO projects VALUES (1, 'Project Alpha', 1, '2023-01-01', 'Active');
-INSERT INTO projects VALUES (2, 'Project Beta', 2, '2023-06-15', 'Completed');
-
-INSERT INTO project_assignments VALUES (1, 1, 100);
-INSERT INTO project_assignments VALUES (2, 1, 150);
-INSERT INTO project_assignments VALUES (4, 2, 80);
+INSERT INTO departments VALUES (1,'Engineering',500000,'New York'), (2,'Marketing',200000,'Chicago'), (3,'Sales',300000,'Los Angeles');
+INSERT INTO employees VALUES (1,'Alice Chen','Engineering',95000,'2020-01-15',NULL), (2,'Bob Smith','Engineering',85000,'2021-03-20',1), (3,'Carol Davis','Marketing',72000,'2019-07-01',NULL), (4,'David Lee','Sales',68000,'2022-11-01',NULL), (5,'Eve Turner','Engineering',110000,'2018-05-10',1), (6,'Frank White','Marketing',65000,'2023-01-15',3);
+INSERT INTO projects VALUES (1,'Data Platform',1,'2023-01-01','active'), (2,'Website Redesign',2,'2023-03-15','completed'), (3,'CRM Migration',3,'2023-06-01','active');
+INSERT INTO project_assignments VALUES (1,1,120),(2,1,80),(5,1,200), (3,2,160),(6,2,40),(4,3,100);
 """
 
-TASKS = [
-    {
-        "id": "easy_01",
-        "name": "Filter high earners",
-        "difficulty": "easy",
-        "type": "write_query",
-        "grader": True,
-        "description": "Find Engineering employees with salary > 90000, return name and salary ordered by salary DESC.",
-        "expected_rows": 2
-    },
-    {
-        "id": "easy_02",
-        "name": "Count per department",
-        "difficulty": "easy",
-        "type": "write_query",
-        "grader": True,
-        "description": "Count employees per department, return department and count ordered by count DESC.",
-        "expected_rows": 3
-    },
-    {
-        "id": "medium_01",
-        "name": "Fix broken JOIN",
-        "difficulty": "medium",
-        "type": "fix_query",
-        "grader": True,
-        "description": "Fix a broken JOIN query (missing ON keyword, wrong table alias in WHERE). Return name and department name.",
-        "expected_rows": 6
-    },
-    {
-        "id": "medium_02",
-        "name": "Fix GROUP BY",
-        "difficulty": "medium",
-        "type": "fix_query",
-        "grader": True,
-        "description": "Fix wrong GROUP BY (should GROUP BY department not id) finding max salary per department.",
-        "expected_rows": 3
-    },
-    {
-        "id": "hard_01",
-        "name": "Optimize correlated subquery",
-        "difficulty": "hard",
-        "type": "optimize_query",
-        "grader": True,
-        "description": "Replace correlated subquery with window function/CTE to find top earner per department with total hours.",
-        "expected_rows": 3
-    },
-    {
-        "id": "hard_02",
-        "name": "Eliminate N+1 subqueries",
-        "difficulty": "hard",
-        "type": "optimize_query",
-        "grader": True,
-        "description": "Eliminate N+1 subqueries using LEFT JOIN + GROUP BY to compute total hours per project.",
-        "expected_rows": 2
-    }
-]
+TASKS = {
+  "easy_01": {
+    "type": "write_query",
+    "difficulty": "easy",
+    "description": "Find all employees in the Engineering department with salary above 90000. Return name and salary ordered by salary descending.",
+    "expected_columns": ["name", "salary"],
+    "expected_row_count": 2,
+    "hints": ["Filter department = Engineering", "salary > 90000", "ORDER BY salary DESC"]
+  },
+  "easy_02": {
+    "type": "write_query", 
+    "difficulty": "easy",
+    "description": "Count how many employees are in each department. Return department and employee count ordered by count descending.",
+    "expected_columns": ["department", "count"],
+    "expected_row_count": 3,
+    "hints": ["Use GROUP BY department", "Use COUNT(*)"]
+  },
+  "medium_01": {
+    "type": "fix_query",
+    "difficulty": "medium",
+    "description": "Fix this broken query that returns employees on active projects with project name and hours worked.",
+    "broken_sql": "SELECT e.name, p.name, pa.hours_worked FROM employees e JOIN project_assignments pa e.id = pa.employee_id JOIN projects p ON pa.project_id = p.id WHERE projects.status = 'active'",
+    "expected_row_count": 3,
+    "bugs": ["Missing ON keyword in first JOIN", "Wrong alias 'projects' should be 'p' in WHERE"]
+  },
+  "medium_02": {
+    "type": "fix_query",
+    "difficulty": "medium", 
+    "description": "Fix this query that should return average salary per department but groups incorrectly.",
+    "broken_sql": "SELECT department, AVG(salary) FROM employees GROUP BY id ORDER BY AVG(salary) DESC",
+    "expected_row_count": 3,
+    "bugs": ["GROUP BY should use department not id"]
+  },
+  "hard_01": {
+    "type": "optimize_query",
+    "difficulty": "hard",
+    "description": "Optimize this slow query: find top earner in each department with their total hours worked. Replace correlated subqueries with window functions or CTEs.",
+    "slow_sql": "SELECT e.name, e.department, e.salary, (SELECT SUM(hours_worked) FROM project_assignments pa WHERE pa.employee_id = e.id) as total_hours FROM employees e WHERE e.salary = (SELECT MAX(salary) FROM employees e2 WHERE e2.department = e.department) ORDER BY e.salary DESC",
+    "expected_row_count": 3,
+    "good_patterns": ["ROW_NUMBER", "RANK", "WITH ", "LEFT JOIN"],
+    "optimization_hints": ["Use window functions", "Use CTEs"]
+  },
+  "hard_02": {
+    "type": "optimize_query",
+    "difficulty": "hard",
+    "description": "Rewrite this N+1 query: return each department name, budget, employee count, and active project count using JOINs instead of subqueries.",
+    "slow_sql": "SELECT d.name, d.budget, (SELECT COUNT(*) FROM employees e WHERE e.department = d.name) as emp_count, (SELECT COUNT(*) FROM projects p WHERE p.department_id = d.id AND p.status='active') as active_projects FROM departments d",
+    "expected_row_count": 3,
+    "good_patterns": ["LEFT JOIN", "GROUP BY", "COUNT("],
+    "optimization_hints": ["Use LEFT JOIN with GROUP BY"]
+  }
+}
 
-def build_db() -> sqlite3.Connection:
-    conn = sqlite3.connect(':memory:')
-    conn.row_factory = sqlite3.Row
-    conn.executescript(SCHEMA_SQL)
-    conn.executescript(SEED_DATA_SQL)
-    conn.commit()
-    return conn
-
-class SqlEnvironment:
+class SQLEnvironment:
     def __init__(self):
-        self.db = build_db()
-        self.max_steps = 8
         self.step_count = 0
-        self.current_task = TASKS[0]
         self.done = False
+        self.reward = 0.05
+        self.feedback = ""
         self.last_sql = ""
         self.last_result = ""
         self.last_error = ""
-        self.reward = 0.05
-        self.feedback = ""
+        self.task_id = "easy_01"
+        self.task = TASKS[self.task_id]
+        self.conn = None
+        self._build_db()
+
+    def _build_db(self):
+        if self.conn:
+            self.conn.close()
+        self.conn = sqlite3.connect(":memory:")
+        self.conn.row_factory = sqlite3.Row
+        self.conn.executescript(SCHEMA_SQL)
+        self.conn.executescript(SEED_DATA_SQL)
+        self.conn.commit()
+        return self.conn
 
     def _clamp(self, score: float) -> float:
         return round(max(0.05, min(0.95, score)), 4)
 
-    def _execute_sql(self, sql: str) -> tuple[Optional[List[Dict[str, Any]]], Optional[str]]:
+    def reset(self, task_id=None, difficulty=None):
+        self.step_count = 0
+        self.done = False
+        self.last_sql = ""
+        self.last_result = ""
+        self.last_error = ""
+        self.reward = self._clamp(0.15)
+        self.feedback = "Environment reset."
+
+        if task_id and task_id in TASKS:
+            self.task_id = task_id
+            self.task = TASKS[task_id]
+        elif difficulty:
+            candidates = [k for k, v in TASKS.items() if v["difficulty"] == difficulty]
+            if candidates:
+                self.task_id = candidates[0]
+                self.task = TASKS[self.task_id]
+        else:
+            self.task_id = list(TASKS.keys())[0]
+            self.task = TASKS[self.task_id]
+
+        self.conn = self._build_db()
+
+        if self.task["type"] == "fix_query":
+            self.last_sql = self.task["broken_sql"]
+        elif self.task["type"] == "optimize_query":
+            self.last_sql = self.task["slow_sql"]
+
+        return self._build_observation(self.reward, self.feedback)
+
+    def _execute_sql(self, sql: str):
         upper_sql = sql.upper()
         blocked = ["DROP", "DELETE", "INSERT", "UPDATE", "CREATE", "ALTER", "TRUNCATE"]
         for b in blocked:
             if re.search(rf"\b{b}\b", upper_sql):
-                return None, f"Blocked keyword {b} found. Only SELECT allowed."
+                return None, "Forbidden operation"
+                
         try:
-            cursor = self.db.cursor()
+            cursor = self.conn.cursor()
             cursor.execute(sql)
             rows = cursor.fetchall()
-            return [dict(row) for row in rows], None
+            if not rows:
+                return json.dumps([{}]) if cursor.description else "[]", None
+            
+            res_dict = [dict(r) for r in rows[:10]]
+            return json.dumps(res_dict), None
         except Exception as e:
             return None, str(e)
 
-    def _grade_write(self, sql: str, rows: Optional[List[Dict]], expected_rows: int) -> float:
+    def _grade_write(self, sql, rows, cols) -> tuple[float, str]:
         score = 0.0
-        upper_sql = sql.upper()
-        if "EMPLOYEES" in upper_sql or "DEPARTMENTS" in upper_sql:
+        feedback_parts = []
+        lower_sql = sql.lower()
+        
+        if "employees" in lower_sql or "departments" in lower_sql:
             score += 0.15
-        if "WHERE" in upper_sql:
+            feedback_parts.append("Correct table used")
+            
+        if "where" in lower_sql:
             score += 0.10
-        if "ORDER BY" in upper_sql:
-            score += 0.15
+            feedback_parts.append("WHERE clause present")
+            
+        expected_cols = [c.lower() for c in self.task["expected_columns"]]
+        actual_cols = [c.lower() for c in cols]
+        matched = sum(1 for ec in expected_cols if ec in actual_cols)
+        if expected_cols:
+            score += (matched / len(expected_cols)) * 0.25
+            feedback_parts.append(f"Columns: {matched}/{len(expected_cols)} matched")
+            
+        expected_rows = self.task["expected_row_count"]
         if rows is not None:
-            if len(rows) > 0:
-                score += 0.25
             if len(rows) == expected_rows:
                 score += 0.25
+                feedback_parts.append(f"Row count correct: {len(rows)}")
             elif abs(len(rows) - expected_rows) == 1:
                 score += 0.12
-        return self._clamp(score)
+                feedback_parts.append(f"Row count close: {len(rows)} vs {expected_rows}")
+            else:
+                feedback_parts.append(f"Row count wrong: {len(rows)} vs {expected_rows}")
+                
+        if "order by" in lower_sql:
+            score += 0.15
+            feedback_parts.append("ORDER BY present")
+            
+        return self._clamp(score), " | ".join(feedback_parts)
 
-    def _grade_fix(self, sql: str, rows: Optional[List[Dict]], expected_rows: int) -> float:
+    def _grade_fix(self, sql, rows, error) -> tuple[float, str]:
         score = 0.0
-        upper_sql = sql.upper()
-        if rows is not None:
+        feedback_parts = []
+        lower_sql = sql.lower()
+        
+        if error is None and rows is not None:
             score += 0.35
-            if len(rows) == expected_rows:
-                score += 0.35
-        if re.search(r"\bJOIN\b.*\bON\b", upper_sql):
+            feedback_parts.append("Query executes successfully")
+        else:
+            feedback_parts.append(f"Still broken: {error}")
+            
+        expected_rows = self.task["expected_row_count"]
+        if rows is not None and len(rows) == expected_rows:
+            score += 0.35
+            feedback_parts.append(f"Correct rows: {len(rows)}")
+            
+        if "join" in lower_sql and " on " in lower_sql:
             score += 0.15
-        if "WHERE" in upper_sql or "GROUP BY" in upper_sql:
+            feedback_parts.append("JOIN...ON syntax correct")
+            
+        if "group by" in lower_sql or "where" in lower_sql:
             score += 0.10
-        return self._clamp(score)
+            feedback_parts.append("Filter/grouping present")
+            
+        return self._clamp(score), " | ".join(feedback_parts)
 
-    def _grade_optimize(self, sql: str, rows: Optional[List[Dict]], expected_rows: int) -> float:
+    def _grade_optimize(self, sql, rows) -> tuple[float, str]:
         score = 0.0
-        upper_sql = sql.upper()
-        if re.search(r"\b(WITH|ROW_NUMBER|RANK)\b", upper_sql):
+        feedback_parts = []
+        sql_upper = sql.upper()
+        
+        if "WITH " in sql_upper or "ROW_NUMBER" in sql_upper or "RANK" in sql_upper:
             score += 0.30
-        if "JOIN" in upper_sql:
+            feedback_parts.append("Uses CTE or window function")
+            
+        if "JOIN" in sql_upper:
             score += 0.25
-        if upper_sql.count("SELECT") == 1:
+            feedback_parts.append("Uses JOIN")
+            
+        select_count = sql_upper.count("SELECT")
+        if select_count == 1:
             score += 0.20
+            feedback_parts.append("No nested SELECT (no correlated subqueries)")
+        else:
+            feedback_parts.append(f"Still has {select_count - 1} nested SELECT(s)")
+            
+        expected_rows = self.task["expected_row_count"]
         if rows is not None and len(rows) == expected_rows:
             score += 0.20
-        return self._clamp(score)
+            feedback_parts.append(f"Correct rows: {len(rows)}")
+            
+        return self._clamp(score), " | ".join(feedback_parts)
 
-    def reset(self, task_id: Optional[str] = None, difficulty: Optional[str] = None) -> Dict[str, Any]:
-        self.step_count = 0
-        self.done = False
-        self.last_sql = ""
-        self.last_result = ""
-        self.last_error = ""
-        self.reward = self._clamp(0.0)
-        self.feedback = "Environment reset."
+    def get_schema_info(self) -> str:
+        return SCHEMA_SQL.strip()
         
-        if task_id:
-            for t in TASKS:
-                if t["id"] == task_id:
-                    self.current_task = t
-                    break
-        elif difficulty:
-            for t in TASKS:
-                if t["difficulty"] == difficulty:
-                    self.current_task = t
-                    break
-        else:
-            self.current_task = TASKS[0]
-            
-        return self.get_observation()
+    def get_sample_data(self) -> str:
+        samples = []
+        tables = ["departments", "employees", "projects", "project_assignments"]
+        cursor = self.conn.cursor()
+        for t in tables:
+            cursor.execute(f"SELECT * FROM {t} LIMIT 3")
+            rows = cursor.fetchall()
+            cols = [desc[0] for desc in cursor.description] if cursor.description else []
+            samples.append(f"Table {t}:\nCols: {', '.join(cols)}")
+            for r in rows:
+                samples.append(str(dict(r)))
+        return "\n".join(samples)
 
-    def get_observation(self) -> Dict[str, Any]:
-        return {
-            "task_id": self.current_task["id"],
-            "task_type": self.current_task["type"],
-            "task_description": self.current_task["description"],
-            "schema_info": SCHEMA_SQL.strip(),
-            "last_sql": self.last_sql,
-            "last_result": self.last_result,
-            "last_error": self.last_error,
+    def _build_observation(self, reward, feedback, last_sql=None, last_result=None, last_error=None):
+        reward = self._clamp(reward)
+        obs = {
+            "task_id": self.task_id,
+            "task_type": self.task["type"],
+            "task_description": self.task["description"],
+            "schema_info": self.get_schema_info(),
+            "sample_data": self.get_sample_data(),
+            "last_sql": last_sql or self.last_sql,
+            "last_result": last_result or self.last_result,
+            "last_error": last_error or self.last_error,
             "step_count": self.step_count,
             "done": self.done,
-            "reward": self.reward,
-            "feedback": self.feedback
+            "reward": reward,
+            "feedback": feedback,
+            "metadata": self.task.get("hints", []) + self.task.get("good_patterns", [])
         }
+        return {"observation": obs, "reward": reward, "done": self.done, "info": {}}
 
     @property
-    def state(self) -> Dict[str, Any]:
+    def state(self):
         return {
-            "task": self.current_task["id"],
-            "step": self.step_count,
-            "done": self.done,
-            "reward": self.reward
+            "episode_id": "ep_" + str(self.step_count),
+            "task_type": self.task["type"],
+            "task_id": self.task_id,
+            "step_count": self.step_count,
+            "total_reward": self.reward,
+            "best_score_so_far": self.reward,
+            "attempts": self.step_count
         }
 
-    def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
-        if self.done:
-            return self.get_observation()
-            
+    def step(self, action: dict):
         self.step_count += 1
         sql = action.get("sql", "").strip()
         self.last_sql = sql
         
         if not sql:
-            self.last_error = "No SQL provided."
-            self.last_result = ""
             self.reward = self._clamp(0.0)
-            self.feedback = "Empty query."
+            self.last_error = "No SQL provided"
+            self.last_result = ""
+            self.feedback = "Missing SQL"
         else:
-            rows, err = self._execute_sql(sql)
+            res_text, err = self._execute_sql(sql)
             if err:
                 self.last_error = err
                 self.last_result = ""
-                self.reward = self._clamp(0.0)
-                self.feedback = "Execution failed."
+                self.reward = self._clamp(0.08)
+                self.feedback = err
             else:
                 self.last_error = ""
-                self.last_result = json.dumps(rows)
-                expected = self.current_task["expected_rows"]
-                ttype = self.current_task["type"]
+                self.last_result = res_text
                 
-                raw_score = 0.0
+                rows = []
+                cols = []
+                if res_text and res_text != "[]" and res_text != "[{}]":
+                    try:
+                        rows = json.loads(res_text)
+                        if rows and isinstance(rows[0], dict):
+                            cols = list(rows[0].keys())
+                    except:
+                        pass
+                
+                ttype = self.task["type"]
                 if ttype == "write_query":
-                    raw_score = self._grade_write(sql, rows, expected)
+                    self.reward, self.feedback = self._grade_write(sql, rows, cols)
                 elif ttype == "fix_query":
-                    raw_score = self._grade_fix(sql, rows, expected)
+                    self.reward, self.feedback = self._grade_fix(sql, rows, err)
                 elif ttype == "optimize_query":
-                    raw_score = self._grade_optimize(sql, rows, expected)
-                
-                self.reward = raw_score # already clamped internally by the functions
-                self.feedback = "Query executed successfully."
+                    self.reward, self.feedback = self._grade_optimize(sql, rows)
 
-        if self.step_count >= self.max_steps:
+        if self.reward > 0.90 or self.step_count >= 8:
             self.done = True
-            
-        return self.get_observation()
+
+        return self._build_observation(self.reward, self.feedback)