File size: 2,780 Bytes
5d911cb 39419fc 5d911cb 39419fc 5d911cb 39419fc 5d911cb 39419fc 5d911cb 39419fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from dataclasses import dataclass, field
from typing import List, Callable, Any
import random
@dataclass
class Task:
id: str
description: str
expected_columns: List[str]
validation_query: str
difficulty: str
grader: Callable[[List[Any]], float] = None # ← add this
def grader_salary_above_50000(rows: List[Any]) -> float:
"""All returned rows must have salary > 50000."""
if not rows:
return 0.0
correct = all(row[1] > 50000 for row in rows)
return 1.0 if correct else 0.0
def grader_employee_department_join(rows: List[Any]) -> float:
"""Must return rows with at least 2 columns (name + department)."""
if not rows:
return 0.0
correct = all(len(row) >= 2 for row in rows)
return 1.0 if correct else 0.0
def grader_avg_salary_per_department(rows: List[Any]) -> float:
"""Must return rows with department name and a numeric avg salary."""
if not rows:
return 0.0
try:
correct = all(isinstance(row[1], (int, float)) for row in rows)
return 1.0 if correct else 0.0
except (IndexError, TypeError):
return 0.0
TASKS = [
Task(
id="task_01",
description="Get the name and salary of all employees earning more than 50000",
expected_columns=["name", "salary"],
validation_query="SELECT name, salary FROM employees WHERE salary > 50000",
difficulty="easy",
grader=grader_salary_above_50000,
),
Task(
id="task_02",
description="Get each employee's name along with their department name",
expected_columns=["name", "department"],
validation_query="SELECT e.name, d.name AS department FROM employees e JOIN departments d ON e.department = d.id",
difficulty="medium",
grader=grader_employee_department_join,
),
Task(
id="task_03",
description="Get the average salary per department with department name",
expected_columns=["name", "avg_salary"],
validation_query="SELECT d.name, AVG(e.salary) AS avg_salary FROM employees e JOIN departments d ON e.department = d.id GROUP BY d.id",
difficulty="hard",
grader=grader_avg_salary_per_department,
),
]
def get_task(task_id: str) -> Task:
for t in TASKS:
if t.id == task_id:
return t
raise ValueError(f"Task {task_id} not found")
def get_random_task() -> Task:
return random.choice(TASKS)
def get_task_by_difficulty(difficulty: str) -> Task:
filtered = [t for t in TASKS if t.difficulty == difficulty]
if not filtered:
raise ValueError(f"No tasks found for difficulty: {difficulty}")
return random.choice(filtered) |