ritvik360/nl2sql-bench

1"""2nl2sql-bench/server/tasks/base.py3==================================4Abstract base for all NL2SQL tasks and the global task registry.5 6Each task holds a list of (question, ground_truth_sql) pairs.7The environment picks one pair per episode via a deterministic round-robin8so that the same task always cycles through the same question sequence —9this keeps grader results reproducible across runs.10"""11 12from __future__ import annotations13 14import sqlite315from abc import ABC, abstractmethod16from typing import Dict, List, NamedTuple, Tuple, Type17 18 19class TaskExample(NamedTuple):20    question: str21    sql: str22    # Human-readable description of what makes this question that difficulty23    notes: str = ""24 25 26class BaseTask(ABC):27    """Abstract base class for all tasks."""28 29    name: str = ""30    difficulty: str = ""   # easy | medium | hard31    examples: List[TaskExample] = []32 33    def __init__(self) -> None:34        if not self.examples:35            raise ValueError(f"Task {self.name!r} has no examples defined.")36        self._cursor = 0   # round-robin index37 38    def next_example(self) -> TaskExample:39        """Return the next question in round-robin order."""40        example = self.examples[self._cursor % len(self.examples)]41        self._cursor += 142        return example43 44    @classmethod45    def schema_context(cls) -> str:46        """Return a compact schema description for the agent system prompt."""47        return _SCHEMA_CONTEXT48 49    @abstractmethod50    def description(self) -> str:51        """One-sentence description for openenv.yaml."""52 53 54# ── Global schema context string (injected into every observation) ─────────55 56_SCHEMA_CONTEXT = """\57Database: ecommerce  (SQLite, read-only)58 59TABLES60------61categories(id INTEGER PK, name TEXT)62 63products(id INTEGER PK, name TEXT, category_id INTEGER FK→categories.id,64         price REAL, stock_quantity INTEGER)65 66customers(id INTEGER PK, name TEXT, email TEXT, country TEXT,67          tier TEXT ∈ {bronze|silver|gold}, created_at TEXT ISO-8601)68 69orders(id INTEGER PK, customer_id INTEGER FK→customers.id,70       status TEXT ∈ {pending|processing|shipped|delivered|cancelled},71       created_at TEXT ISO-8601, total_amount REAL)72 73order_items(id INTEGER PK, order_id INTEGER FK→orders.id,74            product_id INTEGER FK→products.id,75            quantity INTEGER, unit_price REAL)76 77reviews(id INTEGER PK, product_id INTEGER FK→products.id,78        customer_id INTEGER FK→customers.id,79        rating INTEGER 1-5, created_at TEXT ISO-8601)80 81NOTES82-----83- Date comparisons: use   created_at >= '2024-01-01'   (text ISO sort works)84- SQLite window functions (RANK, DENSE_RANK, ROW_NUMBER, LAG, LEAD) are available85- strftime('%Y-%m', created_at) returns 'YYYY-MM' month strings86- All monetary values are in USD87"""88 89 90# ── Task registry ──────────────────────────────────────────────────────────91 92_REGISTRY: Dict[str, Type[BaseTask]] = {}93 94 95def register(cls: Type[BaseTask]) -> Type[BaseTask]:96    """Class decorator to register a task."""97    _REGISTRY[cls.name] = cls98    return cls99 100 101def get_task(name: str) -> BaseTask:102    if name not in _REGISTRY:103        raise KeyError(f"Unknown task {name!r}. Available: {list(_REGISTRY)}")104    return _REGISTRY[name]()105 106 107def all_task_names() -> List[str]:108    return list(_REGISTRY.keys())109