Update server/tasks.py

server/tasks.py

@@ -1,87 +1,131 @@
-"""Task definitions and router for the validator server."""
+"""Deterministic task definitions for the Eco-LLM routing environment."""
 
 from __future__ import annotations
-from typing import Callable, Optional, Any, Dict
-import sys
-
-def _get_server_app():
-    """Lazy import of server.App."""
-    try:
-        from server.app import App
-        return App()
-    except ImportError as e:
-        print(f"Warning: Could not import server.App: {e}", file=sys.stderr)
-        return None
-
-
-def _get_task_registry():
-    """Lazy import of task registry."""
-    try:
-        from tasks import TASKS as TASKS_DEFS
-        return TASKS_DEFS
-    except ImportError as e:
-        print(f"Warning: Could not import TASKS: {e}", file=sys.stderr)
-        return []
-
-
-def get_task(task_id: str, *, task_def: Optional[dict] = None) -> dict:
-    """Get a task definition by ID."""
-    if task_def:
-        return task_def
-    from tasks import TASK_MAP
-    return TASK_MAP.get(task_id, {})
-
-
-# Define task definitions
-TASKS = [
-    {
-        "id": "easy",
-        "task_id": "easy",
-        "name": "easy",
-        "title": "Single Query Routing",
-        "difficulty": "easy",
-        "description": "Route a single LLM query to the optimal model tier while minimising carbon footprint and latency.",
-        "max_steps": 10,
-        "grader": None,
-        "grader_path": "graders:grade_easy",
-        "graders": [],
-        "grader_paths": [],
-        "reward_range": [0.0, 1.0],
-    },
-    {
-        "id": "medium",
-        "task_id": "medium",
-        "name": "medium",
-        "title": "Multi-Query Episode",
-        "difficulty": "medium",
-        "description": "Route 3 queries of varying complexity.",
-        "max_steps": 20,
-        "grader": None,
-        "grader_path": "graders:grade_medium",
-        "graders": [],
-        "grader_paths": [],
-        "reward_range": [0.0, 1.0],
-    },
-    {
-        "id": "hard",
-        "task_id": "hard",
-        "name": "hard",
-        "title": "Stateful Carbon-Aware Routing",
-        "difficulty": "hard",
-        "description": "5-query episode with caching, KB lookups, cascade strategies, and carbon-aware waiting.",
-        "max_steps": 50,
-        "grader": None,
-        "grader_path": "graders:grade_hard",
-        "graders": [],
-        "grader_paths": [],
-        "reward_range": [0.0, 1.0],
-    },
-]
-
-# Build TASK_MAP
-TASK_MAP = {task["task_id"]: task for task in TASKS}
-
-
-def _get_task_map():
-    """Lazy import of TASK_MAP."""
-    return TASK_MAP
+
+from dataclasses import dataclass
+
+from server.models import ModelChoice
+
+
+ENERGY_COSTS = {
+    ModelChoice.SMALL: 0.1,
+    ModelChoice.MEDIUM: 0.3,
+    ModelChoice.LARGE: 0.6,
+}
+
+LATENCY_COSTS = {
+    ModelChoice.SMALL: 1.0,
+    ModelChoice.MEDIUM: 2.0,
+    ModelChoice.LARGE: 5.0,
+}
+
+MODEL_ORDER = [ModelChoice.SMALL, ModelChoice.MEDIUM, ModelChoice.LARGE]
+MODEL_RANK = {model: index for index, model in enumerate(MODEL_ORDER)}
+CARBON_SCHEDULE = [0.8, 0.6, 0.4, 0.3, 0.5]
+LATENCY_MU = 0.1        
+MAX_CARBON_PENALTY = 0.8
+KB_ENERGY_COST = 0.05
+KB_LATENCY_COST = 1.5
+MAX_STEPS_PER_EPISODE = 50
+
+
+@dataclass(frozen=True)
+class QuerySpec:
+    text: str
+    correct_answer: str
+    minimum_model: ModelChoice
+    kb_available: bool = False
+
+
+@dataclass(frozen=True)
+class TaskSpec:
+    task_id: str
+    difficulty: str
+    queries: tuple[QuerySpec, ...]
+    large_model_penalty: float = 0.0
+    description: str = ""
+
+
+TASKS: dict[str, TaskSpec] = {
+    "task_1": TaskSpec(
+        task_id="task_1",
+        difficulty="easy",
+        queries=(
+            QuerySpec(
+                text="What is the most efficient insulation material for a passive house retrofit?",
+                correct_answer="Vacuum insulated panels",
+                minimum_model=ModelChoice.MEDIUM,
+            ),
+        ),
+        description="Single-query routing with accuracy-only emphasis.",
+    ),
+    "task_2": TaskSpec(
+        task_id="task_2",
+        difficulty="medium",
+        queries=(
+            QuerySpec(
+                text="Summarize the trade-off between quantization and factual accuracy in a customer support bot.",
+                correct_answer="Aggressive quantization lowers energy use but can reduce factual reliability.",
+                minimum_model=ModelChoice.SMALL,
+            ),
+            QuerySpec(
+                text="Explain when speculative decoding improves throughput on medium-sized models.",
+                correct_answer="It helps when draft tokens are likely to be accepted by the target model.",
+                minimum_model=ModelChoice.MEDIUM,
+            ),
+            QuerySpec(
+                text="Design a safety fallback for multilingual medical triage.",
+                correct_answer="Use high-capability review with escalation and verified retrieval before final output.",
+                minimum_model=ModelChoice.LARGE,
+            ),
+        ),
+        large_model_penalty=-0.2,
+        description="Three-query episode with an extra penalty for using the LARGE model.",
+    ),
+    "task_3": TaskSpec(
+        task_id="task_3",
+        difficulty="hard",
+        queries=(
+            QuerySpec(
+                text="Classify whether this FAQ request can be answered from the billing knowledge base.",
+                correct_answer="Yes, use the billing KB article.",
+                minimum_model=ModelChoice.SMALL,
+                kb_available=True,
+            ),
+            QuerySpec(
+                text="Route a long-form reasoning question about carbon-aware serving policy.",
+                correct_answer="Use cascade with escalation to a larger model if smaller ones fail.",
+                minimum_model=ModelChoice.MEDIUM,
+            ),
+            QuerySpec(
+                text="Classify whether this FAQ request can be answered from the billing knowledge base.",
+                correct_answer="Yes, use the billing KB article.",
+                minimum_model=ModelChoice.SMALL,
+                kb_available=True,
+            ),
+            QuerySpec(
+                text="Generate a compliant answer for a finance-regulated enterprise user.",
+                correct_answer="Use the large model with a retrieval-backed compliance layer.",
+                minimum_model=ModelChoice.LARGE,
+            ),
+            QuerySpec(
+                text="Route a long-form reasoning question about carbon-aware serving policy.",
+                correct_answer="Use cascade with escalation to a larger model if smaller ones fail.",
+                minimum_model=ModelChoice.MEDIUM,
+            ),
+        ),
+        description="Stateful five-query episode with repeats, caching, cascading, early exit, and waiting.",
+    ),
+}
+
+TASK_ALIASES: dict[str, str] = {
+    "easy": "task_1",
+    "medium": "task_2",
+    "hard": "task_3",
+}
+
+
+def get_task(task_id: str) -> TaskSpec:
+    """Return a task by id, defaulting to the easy task for unknown ids."""
+    resolved_task_id = TASK_ALIASES.get(task_id, task_id)
+    return TASKS.get(resolved_task_id, TASKS["task_1"])