Initial commit

dataset/easy_cases.json

@@ -0,0 +1,125 @@
+[
+  {
+    "id": "easy_001",
+    "category": "syntax",
+    "description": "Fix the SQL syntax error: missing comma in SELECT clause.",
+    "buggy_query": "SELECT id name email FROM users WHERE active = 1",
+    "fixed_query": "SELECT id, name, email FROM users WHERE active = 1",
+    "error_message": "ERROR: syntax error at or near 'name'",
+    "database_schema": {
+      "users": ["id INT PRIMARY KEY", "name VARCHAR(100)", "email VARCHAR(100)", "active INT"]
+    },
+    "expected_output": [
+      {"id": 1, "name": "Alice", "email": "alice@example.com"},
+      {"id": 2, "name": "Bob",   "email": "bob@example.com"}
+    ],
+    "error_type": "syntax",
+    "error_location": "SELECT clause",
+    "fix_description": "Add missing commas between column names in SELECT clause",
+    "estimated_fix_steps": 2,
+    "partial_credit_hints": {
+      "identifies_location": 0.2,
+      "identifies_error_type": 0.2,
+      "correct_fix": 0.5,
+      "explanation": 0.1
+    }
+  },
+  {
+    "id": "easy_002",
+    "category": "syntax",
+    "description": "Fix the SQL syntax error: missing WHERE keyword.",
+    "buggy_query": "SELECT * FROM orders id = 101",
+    "fixed_query": "SELECT * FROM orders WHERE id = 101",
+    "error_message": "ERROR: syntax error at or near 'id'",
+    "database_schema": {
+      "orders": ["id INT PRIMARY KEY", "user_id INT", "total DECIMAL", "status VARCHAR(20)"]
+    },
+    "expected_output": [
+      {"id": 101, "user_id": 5, "total": 250.00, "status": "shipped"}
+    ],
+    "error_type": "syntax",
+    "error_location": "FROM clause",
+    "fix_description": "Add missing WHERE keyword before the filter condition",
+    "estimated_fix_steps": 2,
+    "partial_credit_hints": {
+      "identifies_location": 0.2,
+      "identifies_error_type": 0.2,
+      "correct_fix": 0.5,
+      "explanation": 0.1
+    }
+  },
+  {
+    "id": "easy_003",
+    "category": "syntax",
+    "description": "Fix the SQL syntax error: unclosed string literal.",
+    "buggy_query": "SELECT * FROM products WHERE category = 'electronics",
+    "fixed_query": "SELECT * FROM products WHERE category = 'electronics'",
+    "error_message": "ERROR: unterminated quoted string at or near \"'electronics\"",
+    "database_schema": {
+      "products": ["id INT PRIMARY KEY", "name VARCHAR(100)", "category VARCHAR(50)", "price DECIMAL"]
+    },
+    "expected_output": [
+      {"id": 3, "name": "Laptop", "category": "electronics", "price": 999.99}
+    ],
+    "error_type": "syntax",
+    "error_location": "WHERE clause string literal",
+    "fix_description": "Close the single quote at the end of 'electronics'",
+    "estimated_fix_steps": 1,
+    "partial_credit_hints": {
+      "identifies_location": 0.2,
+      "identifies_error_type": 0.2,
+      "correct_fix": 0.5,
+      "explanation": 0.1
+    }
+  },
+  {
+    "id": "easy_004",
+    "category": "syntax",
+    "description": "Fix the SQL syntax error: wrong keyword ORDER used instead of ORDER BY.",
+    "buggy_query": "SELECT id, name, price FROM products ORDER price DESC",
+    "fixed_query": "SELECT id, name, price FROM products ORDER BY price DESC",
+    "error_message": "ERROR: syntax error at or near 'price'",
+    "database_schema": {
+      "products": ["id INT PRIMARY KEY", "name VARCHAR(100)", "category VARCHAR(50)", "price DECIMAL"]
+    },
+    "expected_output": [
+      {"id": 7, "name": "TV", "price": 1200.00},
+      {"id": 3, "name": "Laptop", "price": 999.99}
+    ],
+    "error_type": "syntax",
+    "error_location": "ORDER clause",
+    "fix_description": "Replace ORDER with ORDER BY",
+    "estimated_fix_steps": 1,
+    "partial_credit_hints": {
+      "identifies_location": 0.2,
+      "identifies_error_type": 0.2,
+      "correct_fix": 0.5,
+      "explanation": 0.1
+    }
+  },
+  {
+    "id": "easy_005",
+    "category": "syntax",
+    "description": "Fix the SQL syntax error: GROUP used instead of GROUP BY.",
+    "buggy_query": "SELECT department, COUNT(*) FROM employees GROUP department",
+    "fixed_query": "SELECT department, COUNT(*) FROM employees GROUP BY department",
+    "error_message": "ERROR: syntax error at or near 'department'",
+    "database_schema": {
+      "employees": ["id INT PRIMARY KEY", "name VARCHAR(100)", "department VARCHAR(50)", "salary DECIMAL"]
+    },
+    "expected_output": [
+      {"department": "Engineering", "count": 12},
+      {"department": "Marketing",   "count": 8}
+    ],
+    "error_type": "syntax",
+    "error_location": "GROUP clause",
+    "fix_description": "Replace GROUP with GROUP BY",
+    "estimated_fix_steps": 1,
+    "partial_credit_hints": {
+      "identifies_location": 0.2,
+      "identifies_error_type": 0.2,
+      "correct_fix": 0.5,
+      "explanation": 0.1
+    }
+  }
+]

dataset/hard_cases.json

@@ -0,0 +1,172 @@
+[
+  {
+    "id": "hard_001",
+    "category": "performance",
+    "description": "Detect and fix N+1 query pattern: query fetches all users then runs a subquery per user inside SELECT, causing O(n) database hits.",
+    "buggy_query": "SELECT id, name, (SELECT COUNT(*) FROM orders WHERE user_id = u.id) as order_count, (SELECT SUM(total) FROM orders WHERE user_id = u.id) as total_spent FROM users u",
+    "fixed_query": "SELECT u.id, u.name, COUNT(o.id) as order_count, COALESCE(SUM(o.total), 0) as total_spent FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name",
+    "error_message": "No error thrown — query is functionally correct but causes severe performance degradation at scale (O(n) subqueries)",
+    "database_schema": {
+      "users":  ["id INT PRIMARY KEY", "name VARCHAR(100)", "email VARCHAR(100)"],
+      "orders": ["id INT PRIMARY KEY", "user_id INT REFERENCES users(id)", "total DECIMAL", "status VARCHAR(20)"]
+    },
+    "performance_issue": {
+      "type": "N+1 query",
+      "impact": "With 10,000 users: 20,001 database round trips vs 1",
+      "execution_time_buggy_ms": 8500,
+      "execution_time_fixed_ms": 45
+    },
+    "expected_output": [
+      {"id": 1, "name": "Alice", "order_count": 3, "total_spent": 750.00},
+      {"id": 2, "name": "Bob",   "order_count": 0, "total_spent": 0.00}
+    ],
+    "error_type": "performance",
+    "error_location": "Correlated subqueries in SELECT clause",
+    "fix_description": "Replace correlated subqueries with a single LEFT JOIN + GROUP BY aggregation using COALESCE for null safety",
+    "estimated_fix_steps": 6,
+    "scoring_rubric": {
+      "identifies_n_plus_1_pattern": 0.15,
+      "identifies_correlated_subquery_cause": 0.1,
+      "proposes_join_solution": 0.1,
+      "correct_left_join": 0.1,
+      "correct_aggregation": 0.1,
+      "coalesce_null_safety": 0.05,
+      "explanation_quality": 0.1,
+      "confidence": 0.05
+    },
+    "frontier_model_expected_score": 0.15
+  },
+  {
+    "id": "hard_002",
+    "category": "performance",
+    "description": "Fix missing index causing full table scan: query filters on non-indexed column in 10M row table, fix by identifying missing index and rewriting with index-friendly pattern.",
+    "buggy_query": "SELECT * FROM logs WHERE YEAR(created_at) = 2024 AND MONTH(created_at) = 3",
+    "fixed_query": "SELECT * FROM logs WHERE created_at >= '2024-03-01' AND created_at < '2024-04-01'",
+    "error_message": "No error thrown — runs but causes full table scan because function on indexed column prevents index use",
+    "database_schema": {
+      "logs": ["id BIGINT PRIMARY KEY", "user_id INT", "action VARCHAR(100)", "created_at TIMESTAMP", "INDEX(created_at)"]
+    },
+    "performance_issue": {
+      "type": "Function on indexed column prevents index use (full table scan)",
+      "impact": "10M row table: full scan 12s vs index seek 8ms",
+      "execution_time_buggy_ms": 12000,
+      "execution_time_fixed_ms": 8,
+      "rows_scanned_buggy": 10000000,
+      "rows_scanned_fixed": 85000
+    },
+    "expected_output": "Same rows — functionally equivalent but uses index",
+    "error_type": "performance",
+    "error_location": "WHERE clause: YEAR() and MONTH() functions prevent index usage on created_at",
+    "fix_description": "Replace function-wrapped column with range comparison using explicit date literals so the index on created_at is used",
+    "estimated_fix_steps": 5,
+    "scoring_rubric": {
+      "identifies_function_on_column_issue": 0.15,
+      "identifies_index_not_used": 0.15,
+      "correct_range_rewrite": 0.15,
+      "correct_date_boundaries": 0.1,
+      "explanation_quality": 0.1,
+      "confidence": 0.05
+    },
+    "frontier_model_expected_score": 0.12
+  },
+  {
+    "id": "hard_003",
+    "category": "performance",
+    "description": "Fix implicit cartesian product: missing JOIN condition between two large tables causes cross join — billions of rows.",
+    "buggy_query": "SELECT c.name, o.total FROM customers c, orders o WHERE o.total > 1000",
+    "fixed_query": "SELECT c.name, o.total FROM customers c JOIN orders o ON c.id = o.customer_id WHERE o.total > 1000",
+    "error_message": "No error thrown — query returns astronomically wrong result set (cartesian product)",
+    "database_schema": {
+      "customers": ["id INT PRIMARY KEY", "name VARCHAR(100)", "country VARCHAR(50)"],
+      "orders":    ["id INT PRIMARY KEY", "customer_id INT REFERENCES customers(id)", "total DECIMAL", "status VARCHAR(20)"]
+    },
+    "performance_issue": {
+      "type": "Implicit cartesian product (missing JOIN condition)",
+      "impact": "50K customers × 200K orders = 10 billion intermediate rows",
+      "execution_time_buggy_ms": 999999,
+      "execution_time_fixed_ms": 120
+    },
+    "expected_output": [
+      {"name": "Alice Corp", "total": 15000.00},
+      {"name": "Bob Ltd",    "total": 3200.00}
+    ],
+    "error_type": "performance",
+    "error_location": "FROM clause — two tables listed with comma but no JOIN condition",
+    "fix_description": "Replace implicit comma join with explicit JOIN ON c.id = o.customer_id to eliminate cartesian product",
+    "estimated_fix_steps": 5,
+    "scoring_rubric": {
+      "identifies_cartesian_product": 0.15,
+      "identifies_missing_join_condition": 0.15,
+      "correct_explicit_join": 0.15,
+      "correct_join_condition": 0.1,
+      "explanation_quality": 0.1,
+      "confidence": 0.05
+    },
+    "frontier_model_expected_score": 0.18
+  },
+  {
+    "id": "hard_004",
+    "category": "performance",
+    "description": "Fix SELECT * in JOIN query causing unnecessary data transfer and preventing covering index: rewrite to select only needed columns.",
+    "buggy_query": "SELECT * FROM orders o JOIN users u ON o.user_id = u.id JOIN products p ON o.product_id = p.id WHERE o.status = 'pending'",
+    "fixed_query": "SELECT o.id, o.total, o.status, u.name as user_name, u.email, p.name as product_name, p.price FROM orders o JOIN users u ON o.user_id = u.id JOIN products p ON o.product_id = p.id WHERE o.status = 'pending'",
+    "error_message": "No error — functionally correct but causes over-fetching, prevents covering index, increases memory pressure",
+    "database_schema": {
+      "orders":   ["id INT PRIMARY KEY", "user_id INT", "product_id INT", "total DECIMAL", "status VARCHAR(20)", "created_at TIMESTAMP", "updated_at TIMESTAMP", "notes TEXT"],
+      "users":    ["id INT PRIMARY KEY", "name VARCHAR(100)", "email VARCHAR(100)", "password_hash VARCHAR(255)", "created_at TIMESTAMP"],
+      "products": ["id INT PRIMARY KEY", "name VARCHAR(100)", "price DECIMAL", "stock INT", "description TEXT", "image_url VARCHAR(500)"]
+    },
+    "performance_issue": {
+      "type": "SELECT * with JOINs causes over-fetching and prevents covering index",
+      "impact": "Fetches 15+ columns per row including TEXT blobs vs 7 needed columns. 3x memory overhead.",
+      "execution_time_buggy_ms": 890,
+      "execution_time_fixed_ms": 210
+    },
+    "error_type": "performance",
+    "error_location": "SELECT * across 3-table JOIN",
+    "fix_description": "Replace SELECT * with explicit column list selecting only the 7 columns actually needed by the application",
+    "estimated_fix_steps": 5,
+    "scoring_rubric": {
+      "identifies_select_star_issue": 0.15,
+      "identifies_over_fetching": 0.1,
+      "identifies_covering_index_benefit": 0.1,
+      "correct_column_selection": 0.15,
+      "correct_table_aliases": 0.1,
+      "explanation_quality": 0.1,
+      "confidence": 0.05
+    },
+    "frontier_model_expected_score": 0.20
+  },
+  {
+    "id": "hard_005",
+    "category": "performance",
+    "description": "Fix window function misuse: ROW_NUMBER() applied without PARTITION causing global ranking instead of per-department ranking, and without ORDER BY making results non-deterministic.",
+    "buggy_query": "SELECT id, name, department, salary, ROW_NUMBER() OVER () as rank FROM employees WHERE ROW_NUMBER() OVER () <= 3",
+    "fixed_query": "SELECT id, name, department, salary, rank FROM (SELECT id, name, department, salary, ROW_NUMBER() OVER (PARTITION BY department ORDER BY salary DESC) as rank FROM employees) ranked WHERE rank <= 3",
+    "error_message": "ERROR: window functions are not allowed in WHERE clause",
+    "database_schema": {
+      "employees": ["id INT PRIMARY KEY", "name VARCHAR(100)", "department VARCHAR(50)", "salary DECIMAL", "hire_date DATE"]
+    },
+    "expected_output": [
+      {"id": 3,  "name": "Carol", "department": "Engineering", "salary": 120000, "rank": 1},
+      {"id": 7,  "name": "Dave",  "department": "Engineering", "salary": 110000, "rank": 2},
+      {"id": 12, "name": "Eve",   "department": "Marketing",   "salary": 95000,  "rank": 1}
+    ],
+    "error_type": "performance",
+    "error_location": "Window function in WHERE clause and missing PARTITION BY + ORDER BY",
+    "fix_description": "Wrap in subquery to filter on window function result, add PARTITION BY department and ORDER BY salary DESC for correct per-department ranking",
+    "estimated_fix_steps": 7,
+    "scoring_rubric": {
+      "identifies_window_in_where_error": 0.1,
+      "identifies_missing_partition_by": 0.1,
+      "identifies_missing_order_by": 0.1,
+      "correct_subquery_wrapper": 0.1,
+      "correct_partition_by": 0.1,
+      "correct_order_by": 0.1,
+      "correct_where_on_subquery": 0.1,
+      "explanation_quality": 0.1,
+      "confidence": 0.05
+    },
+    "frontier_model_expected_score": 0.10
+  }
+]

dataset/medium_cases.json

@@ -0,0 +1,136 @@
+[
+  {
+    "id": "medium_001",
+    "category": "logic",
+    "description": "Fix the wrong JOIN type: INNER JOIN excludes users with no orders. Should be LEFT JOIN to include all users.",
+    "buggy_query": "SELECT u.id, u.name, COUNT(o.id) as order_count FROM users u INNER JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name",
+    "fixed_query": "SELECT u.id, u.name, COUNT(o.id) as order_count FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name",
+    "error_message": "No error thrown — query runs but returns wrong results (missing users with 0 orders)",
+    "database_schema": {
+      "users":  ["id INT PRIMARY KEY", "name VARCHAR(100)", "email VARCHAR(100)"],
+      "orders": ["id INT PRIMARY KEY", "user_id INT REFERENCES users(id)", "total DECIMAL", "status VARCHAR(20)"]
+    },
+    "expected_output": [
+      {"id": 1, "name": "Alice", "order_count": 3},
+      {"id": 2, "name": "Bob",   "order_count": 0},
+      {"id": 3, "name": "Carol", "order_count": 1}
+    ],
+    "actual_buggy_output": [
+      {"id": 1, "name": "Alice", "order_count": 3},
+      {"id": 3, "name": "Carol", "order_count": 1}
+    ],
+    "error_type": "logic",
+    "error_location": "JOIN type",
+    "fix_description": "Change INNER JOIN to LEFT JOIN so users with zero orders are included",
+    "estimated_fix_steps": 3,
+    "partial_credit_hints": {
+      "identifies_missing_rows": 0.15,
+      "identifies_join_type_issue": 0.2,
+      "correct_join_fix": 0.4,
+      "explanation": 0.15,
+      "confidence": 0.1
+    }
+  },
+  {
+    "id": "medium_002",
+    "category": "logic",
+    "description": "Fix the wrong JOIN condition: joining on wrong column causes cartesian-like explosion of results.",
+    "buggy_query": "SELECT o.id, o.total, p.name FROM orders o JOIN order_items oi ON o.id = oi.order_id JOIN products p ON oi.id = p.id",
+    "fixed_query": "SELECT o.id, o.total, p.name FROM orders o JOIN order_items oi ON o.id = oi.order_id JOIN products p ON oi.product_id = p.id",
+    "error_message": "No error thrown — query runs but returns incorrect product associations",
+    "database_schema": {
+      "orders":      ["id INT PRIMARY KEY", "user_id INT", "total DECIMAL"],
+      "order_items": ["id INT PRIMARY KEY", "order_id INT", "product_id INT", "quantity INT"],
+      "products":    ["id INT PRIMARY KEY", "name VARCHAR(100)", "price DECIMAL"]
+    },
+    "expected_output": [
+      {"id": 101, "total": 250.00, "name": "Laptop"},
+      {"id": 102, "total": 89.99,  "name": "Mouse"}
+    ],
+    "error_type": "logic",
+    "error_location": "Second JOIN condition — oi.id should be oi.product_id",
+    "fix_description": "Change ON oi.id = p.id to ON oi.product_id = p.id",
+    "estimated_fix_steps": 3,
+    "partial_credit_hints": {
+      "identifies_wrong_column": 0.2,
+      "identifies_join_condition": 0.2,
+      "correct_fix": 0.4,
+      "explanation": 0.15,
+      "confidence": 0.05
+    }
+  },
+  {
+    "id": "medium_003",
+    "category": "logic",
+    "description": "Fix the aggregation logic: HAVING clause filters before GROUP BY aggregation is complete, wrong column used.",
+    "buggy_query": "SELECT department, AVG(salary) as avg_salary FROM employees WHERE AVG(salary) > 70000 GROUP BY department",
+    "fixed_query": "SELECT department, AVG(salary) as avg_salary FROM employees GROUP BY department HAVING AVG(salary) > 70000",
+    "error_message": "ERROR: aggregate functions are not allowed in WHERE",
+    "database_schema": {
+      "employees": ["id INT PRIMARY KEY", "name VARCHAR(100)", "department VARCHAR(50)", "salary DECIMAL"]
+    },
+    "expected_output": [
+      {"department": "Engineering", "avg_salary": 95000.00},
+      {"department": "Data Science","avg_salary": 88000.00}
+    ],
+    "error_type": "logic",
+    "error_location": "WHERE clause — aggregate function used in WHERE instead of HAVING",
+    "fix_description": "Move AVG(salary) > 70000 from WHERE to HAVING clause, placed after GROUP BY",
+    "estimated_fix_steps": 3,
+    "partial_credit_hints": {
+      "identifies_where_vs_having": 0.2,
+      "identifies_aggregate_misuse": 0.2,
+      "correct_fix": 0.4,
+      "explanation": 0.15,
+      "confidence": 0.05
+    }
+  },
+  {
+    "id": "medium_004",
+    "category": "logic",
+    "description": "Fix the subquery logic: correlated subquery compares wrong column, returning incorrect filtered set.",
+    "buggy_query": "SELECT id, name, salary FROM employees e WHERE salary > (SELECT AVG(salary) FROM employees WHERE department = e.id)",
+    "fixed_query": "SELECT id, name, salary FROM employees e WHERE salary > (SELECT AVG(salary) FROM employees WHERE department = e.department)",
+    "error_message": "No error thrown — returns incorrect rows because subquery correlates on wrong column",
+    "database_schema": {
+      "employees": ["id INT PRIMARY KEY", "name VARCHAR(100)", "department VARCHAR(50)", "salary DECIMAL"]
+    },
+    "expected_output": [
+      {"id": 3, "name": "Carol", "salary": 110000},
+      {"id": 7, "name": "Dave",  "salary": 98000}
+    ],
+    "error_type": "logic",
+    "error_location": "Correlated subquery WHERE clause — e.id should be e.department",
+    "fix_description": "Change WHERE department = e.id to WHERE department = e.department",
+    "estimated_fix_steps": 4,
+    "partial_credit_hints": {
+      "identifies_correlated_subquery": 0.15,
+      "identifies_wrong_correlation": 0.2,
+      "correct_fix": 0.45,
+      "explanation": 0.15,
+      "confidence": 0.05
+    }
+  },
+  {
+    "id": "medium_005",
+    "category": "logic",
+    "description": "Fix the DISTINCT misuse: DISTINCT applied to wrong scope, counting duplicates incorrectly.",
+    "buggy_query": "SELECT COUNT(DISTINCT *) FROM orders WHERE status = 'completed'",
+    "fixed_query": "SELECT COUNT(DISTINCT id) FROM orders WHERE status = 'completed'",
+    "error_message": "ERROR: COUNT(DISTINCT) with multiple columns requires explicit column list",
+    "database_schema": {
+      "orders": ["id INT PRIMARY KEY", "user_id INT", "total DECIMAL", "status VARCHAR(20)"]
+    },
+    "expected_output": [{"count": 47}],
+    "error_type": "logic",
+    "error_location": "COUNT(DISTINCT *) — wildcard not valid with DISTINCT",
+    "fix_description": "Replace COUNT(DISTINCT *) with COUNT(DISTINCT id) to count unique order IDs",
+    "estimated_fix_steps": 2,
+    "partial_credit_hints": {
+      "identifies_distinct_misuse": 0.25,
+      "correct_fix": 0.5,
+      "explanation": 0.15,
+      "confidence": 0.1
+    }
+  }
+]

env/models.py

@@ -0,0 +1,195 @@
+from pydantic import BaseModel, Field, field_validator
+from typing import Optional, Any
+from enum import Enum
+import time
+
+
+#  ENUMS
+
+class DifficultyLevel(str, Enum):
+    EASY   = "easy"
+    MEDIUM = "medium"
+    HARD   = "hard"
+
+
+class ActionType(str, Enum):
+    IDENTIFY_ERROR    = "identify_error"
+    PROPOSE_FIX       = "propose_fix"
+    SUBMIT_ANSWER     = "submit_answer"
+    REQUEST_HINT      = "request_hint"
+    EXPLAIN_ISSUE     = "explain_issue"
+    OPTIMIZE_QUERY    = "optimize_query"
+
+# CORE MODELS
+
+
+class Observation(BaseModel):
+    task_id:          str            = Field(..., description="Unique task identifier")
+    task_description: str            = Field(..., description="What the agent must do")
+    current_context:  dict           = Field(..., description="What the agent currently sees")
+    step_count:       int            = Field(default=0, ge=0, description="Steps taken so far")
+    difficulty:       DifficultyLevel = Field(..., description="Task difficulty level")
+    max_steps:        int            = Field(default=20, description="Maximum steps allowed")
+    hints_used:       int            = Field(default=0, description="Number of hints used")
+    previous_actions: list[str]      = Field(default_factory=list, description="History of action types taken")
+    metadata:         dict           = Field(default_factory=dict, description="Extra task metadata")
+
+    model_config = {"json_schema_extra": {
+        "example": {
+            "task_id": "easy_001",
+            "task_description": "Fix the SQL syntax error in the query below.",
+            "current_context": {
+                "buggy_query": "SELECT id, name FROM users WHERE id = 1 AND",
+                "error_message": "SyntaxError: unexpected end of input",
+                "database_schema": "users(id INT, name VARCHAR, email VARCHAR)"
+            },
+            "step_count": 0,
+            "difficulty": "easy",
+            "max_steps": 20,
+            "hints_used": 0,
+            "previous_actions": [],
+            "metadata": {"category": "syntax", "estimated_fix_steps": 2}
+        }
+    }}
+
+
+class Action(BaseModel):
+    action_type: ActionType = Field(..., description="Type of action the agent is taking")
+    payload:     dict       = Field(..., description="Action-specific data")
+
+    @field_validator("payload")
+    @classmethod
+    def payload_must_not_be_empty(cls, v):
+        if v is None:
+            raise ValueError("Payload cannot be None")
+        return v
+
+    @field_validator("payload")
+    @classmethod
+    def truncate_long_strings(cls, v):
+        # Edge case: extremely long agent output — truncate gracefully
+        def truncate(obj, max_len=5000):
+            if isinstance(obj, str) and len(obj) > max_len:
+                return obj[:max_len] + "...[truncated]"
+            if isinstance(obj, dict):
+                return {k: truncate(val, max_len) for k, val in obj.items()}
+            return obj
+        return truncate(v)
+
+    model_config = {"json_schema_extra": {
+        "example": {
+            "action_type": "submit_answer",
+            "payload": {
+                "fixed_query":   "SELECT id, name FROM users WHERE id = 1",
+                "explanation":   "Removed the trailing AND which caused a syntax error",
+                "error_type":    "syntax",
+                "confidence":    0.95
+            }
+        }
+    }}
+
+
+class Reward(BaseModel):
+    score:     float = Field(..., ge=-1.0, le=1.0, description="Reward score between -1.0 and 1.0")
+    breakdown: dict  = Field(..., description="Partial credit details per dimension")
+    feedback:  str   = Field(..., description="Human-readable explanation of the reward")
+
+    @field_validator("score")
+    @classmethod
+    def clamp_score(cls, v):
+        return max(-1.0, min(1.0, round(v, 4)))
+
+    model_config = {"json_schema_extra": {
+        "example": {
+            "score": 0.75,
+            "breakdown": {
+                "correct_answer":  0.5,
+                "explanation":     0.2,
+                "confidence":      0.05,
+                "step_efficiency": 0.0
+            },
+            "feedback": "Correct fix applied. Good explanation provided. Minor efficiency penalty."
+        }
+    }}
+
+
+#  EPISODE STATE (used by state() endpoint)
+
+class EpisodeState(BaseModel):
+    task_id:          Optional[str]            = Field(default=None)
+    difficulty:       Optional[DifficultyLevel] = Field(default=None)
+    step_count:       int                       = Field(default=0)
+    total_reward:     float                     = Field(default=0.0)
+    done:             bool                      = Field(default=False)
+    hints_used:       int                       = Field(default=0)
+    previous_actions: list[str]                 = Field(default_factory=list)
+    action_counts:    dict[str, int]            = Field(default_factory=dict)
+    started_at:       Optional[float]           = Field(default=None)
+    last_reward:      float                     = Field(default=0.0)
+    initialized:      bool                      = Field(default=False)
+
+    model_config = {"json_schema_extra": {
+        "example": {
+            "task_id":          "medium_002",
+            "difficulty":       "medium",
+            "step_count":       3,
+            "total_reward":     0.45,
+            "done":             False,
+            "hints_used":       1,
+            "previous_actions": ["identify_error", "request_hint", "propose_fix"],
+            "action_counts":    {"identify_error": 1, "request_hint": 1, "propose_fix": 1},
+            "started_at":       1700000000.0,
+            "last_reward":      0.25,
+            "initialized":      True
+        }
+    }}
+
+
+#  API REQUEST / RESPONSE WRAPPERS
+
+class StepResponse(BaseModel):
+    observation: Observation
+    reward:      Reward
+    done:        bool
+    info:        dict
+
+class ResetResponse(BaseModel):
+    observation: Observation
+
+class TaskInfo(BaseModel):
+    id:           str
+    difficulty:   DifficultyLevel
+    description:  str
+    action_schema: dict   # REQUIRED by validator — field definitions not just names
+
+class TaskListResponse(BaseModel):
+    tasks:         list[TaskInfo]
+    total:         int
+    action_types:  list[str]
+
+class BaselineResult(BaseModel):
+    task_id:    str
+    difficulty: DifficultyLevel
+    score:      float
+    steps:      int
+    feedback:   str
+
+class BaselineResponse(BaseModel):
+    results:      list[BaselineResult]
+    average_score: float
+    completed_at:  float = Field(default_factory=time.time)
+
+class GraderRequest(BaseModel):
+    task_id:  str
+    action:   Optional[Action] = None
+    episode:  Optional[dict]   = None
+
+class GraderResponse(BaseModel):
+    score:    float
+    feedback: str
+    breakdown: dict
+
+class HealthResponse(BaseModel):
+    status:  str = "ok"
+    version: str = "1.0.0"
+    uptime:  float = Field(default_factory=time.time)

env/tasks.py

@@ -0,0 +1,183 @@
+import json
+import random
+from pathlib import Path
+from env.models import DifficultyLevel, TaskInfo
+
+#  LOAD DATASETS
+
+BASE_DIR = Path(__file__).parent.parent / "dataset"
+
+def _load(filename: str) -> list[dict]:
+    path = BASE_DIR / filename
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+EASY_CASES   = _load("easy_cases.json")
+MEDIUM_CASES = _load("medium_cases.json")
+HARD_CASES   = _load("hard_cases.json")
+
+ALL_CASES: dict[str, list[dict]] = {
+    DifficultyLevel.EASY:   EASY_CASES,
+    DifficultyLevel.MEDIUM: MEDIUM_CASES,
+    DifficultyLevel.HARD:   HARD_CASES,
+}
+
+#  ACTION SCHEMA (required by /tasks validator)
+
+ACTION_SCHEMA = {
+    "identify_error": {
+        "description": "Identify where and what the error is without fixing it yet",
+        "payload_fields": {
+            "error_location": {"type": "string", "required": True,  "description": "Where in the query the error occurs"},
+            "error_type":     {"type": "string", "required": True,  "description": "Type: syntax | logic | performance"},
+            "explanation":    {"type": "string", "required": False, "description": "Brief explanation of the error"}
+        }
+    },
+    "propose_fix": {
+        "description": "Propose a fix without submitting as final answer",
+        "payload_fields": {
+            "fixed_query":  {"type": "string", "required": True,  "description": "The proposed corrected SQL query"},
+            "change_made":  {"type": "string", "required": True,  "description": "What specifically was changed"},
+            "confidence":   {"type": "float",  "required": False, "description": "Confidence score 0.0-1.0"}
+        }
+    },
+    "submit_answer": {
+        "description": "Submit the final fixed query as the definitive answer",
+        "payload_fields": {
+            "fixed_query":   {"type": "string", "required": True,  "description": "Final corrected SQL query"},
+            "explanation":   {"type": "string", "required": True,  "description": "Full explanation of what was wrong and how it was fixed"},
+            "error_type":    {"type": "string", "required": False, "description": "Type: syntax | logic | performance"},
+            "confidence":    {"type": "float",  "required": False, "description": "Confidence score 0.0-1.0"}
+        }
+    },
+    "request_hint": {
+        "description": "Request a hint — costs 0.05 reward penalty per hint",
+        "payload_fields": {
+            "hint_type": {"type": "string", "required": False, "description": "Type of hint wanted: location | error_type | fix_direction"}
+        }
+    },
+    "explain_issue": {
+        "description": "Explain the issue in detail — earns partial credit even without fixing",
+        "payload_fields": {
+            "explanation":    {"type": "string", "required": True,  "description": "Detailed explanation of the SQL problem"},
+            "impact":         {"type": "string", "required": False, "description": "What impact the bug has on query results or performance"},
+            "root_cause":     {"type": "string", "required": False, "description": "Root cause analysis"}
+        }
+    },
+    "optimize_query": {
+        "description": "Submit an optimized version of the query (used for hard/performance tasks)",
+        "payload_fields": {
+            "optimized_query":    {"type": "string", "required": True,  "description": "The performance-optimized SQL query"},
+            "optimization_type":  {"type": "string", "required": True,  "description": "What optimization was applied"},
+            "expected_improvement":{"type": "string", "required": False, "description": "Expected performance gain description"},
+            "explanation":        {"type": "string", "required": False, "description": "Why this optimization works"},
+            "confidence":         {"type": "float",  "required": False, "description": "Confidence 0.0-1.0"}
+        }
+    }
+}
+#  TASK MANAGER
+
+
+class TaskManager:
+    """
+    Manages task selection, hint generation, and task metadata.
+    All tasks are loaded from JSON datasets — no hardcoded tasks.
+    """
+
+    def __init__(self):
+        self._used_ids: set[str] = set()
+
+    def get_task(self, difficulty: DifficultyLevel, task_id: str | None = None) -> dict:
+        """
+        Returns a task dict for the given difficulty.
+        If task_id is provided, returns that specific task.
+        Otherwise picks randomly, avoiding recently used tasks.
+        """
+        pool = ALL_CASES[difficulty]
+
+        if task_id:
+            for case in pool:
+                if case["id"] == task_id:
+                    return case
+            raise ValueError(f"Task '{task_id}' not found in {difficulty} pool")
+
+        # Avoid repeating recently used tasks
+        available = [c for c in pool if c["id"] not in self._used_ids]
+        if not available:
+            self._used_ids.clear()
+            available = pool
+
+        task = random.choice(available)
+        self._used_ids.add(task["id"])
+        return task
+
+    def get_random_task(self) -> dict:
+        """Pick a random task from any difficulty."""
+        difficulty = random.choice(list(DifficultyLevel))
+        return self.get_task(difficulty)
+
+    def build_observation_context(self, task: dict) -> dict:
+        """
+        Build the current_context dict for the Observation.
+        CRITICAL: Must NOT leak the fixed_query (ground truth) to the agent.
+        """
+        context = {
+            "buggy_query":       task["buggy_query"],
+            "error_message":     task["error_message"],
+            "database_schema":   task["database_schema"],
+            "error_type_hint":   task["error_type"],
+            "category":          task["category"],
+            "estimated_steps":   task["estimated_fix_steps"],
+        }
+
+        # For performance tasks include extra context
+        if task.get("performance_issue"):
+            context["performance_issue"] = {
+                "type":   task["performance_issue"]["type"],
+                "impact": task["performance_issue"]["impact"],
+                # Do NOT include timing numbers — agent must figure it out
+            }
+
+        # Include expected output shape (but not the fixed query!)
+        if task.get("expected_output") and isinstance(task["expected_output"], list):
+            context["expected_output_sample"] = task["expected_output"][:1]
+
+        return context
+
+    def get_hint(self, task: dict, hint_number: int) -> str:
+        """
+        Returns progressive hints. Each hint gives more info.
+        Hints cost -0.05 reward each (handled in reward.py).
+        """
+        hints = [
+            f"Hint 1: The error is in the {task.get('error_location', 'query')}.",
+            f"Hint 2: This is a {task.get('error_type', 'unknown')} type error. Category: {task.get('category')}.",
+            f"Hint 3: Fix description — {task.get('fix_description', 'Review the query carefully.')}",
+        ]
+        idx = min(hint_number - 1, len(hints) - 1)
+        return hints[idx]
+
+    def list_all_tasks(self) -> list[TaskInfo]:
+        """Returns TaskInfo list for the /tasks endpoint."""
+        result = []
+        for difficulty, cases in ALL_CASES.items():
+            for case in cases:
+                result.append(TaskInfo(
+                    id=case["id"],
+                    difficulty=difficulty,
+                    description=case["description"],
+                    action_schema=ACTION_SCHEMA
+                ))
+        return result
+
+    def get_ground_truth(self, task_id: str) -> dict | None:
+        """Returns the full ground truth for a task (used by grader only)."""
+        for cases in ALL_CASES.values():
+            for case in cases:
+                if case["id"] == task_id:
+                    return case
+        return None
+
+
+# Singleton instance
+task_manager = TaskManager()