Upload aco/router.py

aco/router.py

@@ -0,0 +1,245 @@
+"""Model Cascade Router - Module 3.
+
+Routes agent requests through a cascade of models:
+tiny local → cheap small cloud → medium → frontier → specialist.
+
+Supports:
+A. always frontier
+B. static routing
+C. prompt-only router
+D. trained cost-aware router
+E. trained router + verifier fallback
+"""
+
+import random
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+from .trace_schema import TaskType, Outcome
+from .config import ACOConfig, ModelConfig
+from .classifier import TaskPrediction
+
+
+@dataclass
+class RoutingDecision:
+    model_id: str
+    provider: str
+    tier: int
+    confidence: float
+    reasoning: str
+    fallback_model_id: Optional[str] = None
+    use_verifier: bool = False
+    max_tokens: int = 4096
+    temperature: float = 0.7
+
+
+class ModelCascadeRouter:
+    """Routes tasks to the cheapest acceptable model."""
+
+    TIER_ORDER = [1, 2, 3, 4, 5]  # tiny → cheap → medium → frontier → specialist
+
+    def __init__(self, config: ACOConfig):
+        self.config = config
+        self.models_by_tier: Dict[int, List[ModelConfig]] = {t: [] for t in self.TIER_ORDER}
+        self._build_tier_index()
+        self.routing_stats: Dict[str, Dict] = {}
+        self.decision_history: List[Dict] = []
+
+    def _build_tier_index(self):
+        for name, mc in self.config.models.items():
+            self.models_by_tier.setdefault(mc.strength_tier, []).append(mc)
+
+    def route(self, task_prediction: TaskPrediction, routing_mode: str = "cascade") -> RoutingDecision:
+        """Select model based on task prediction and routing policy."""
+        
+        if routing_mode == "always_frontier":
+            return self._route_always_frontier(task_prediction)
+        elif routing_mode == "static":
+            return self._route_static(task_prediction)
+        elif routing_mode == "prompt_only":
+            return self._route_prompt_only(task_prediction)
+        elif routing_mode == "learned":
+            return self._route_learned(task_prediction)
+        elif routing_mode == "learned_verifier":
+            return self._route_learned(task_prediction, verifier_fallback=True)
+        else:
+            return self._route_cascade(task_prediction)
+
+    def _route_always_frontier(self, prediction: TaskPrediction) -> RoutingDecision:
+        frontier = self.models_by_tier.get(4, [])
+        if not frontier:
+            frontier = self.models_by_tier.get(5, [])
+        if not frontier:
+            frontier = self.models_by_tier.get(3, [])
+        model = frontier[0] if frontier else list(self.config.models.values())[0]
+        return RoutingDecision(
+            model_id=model.model_id,
+            provider=model.provider,
+            tier=4,
+            confidence=1.0,
+            reasoning="Always frontier policy",
+            max_tokens=min(prediction.expected_cost * 50000, model.max_context),
+        )
+
+    def _route_static(self, prediction: TaskPrediction) -> RoutingDecision:
+        # Static mapping: task type -> tier
+        static_map = {
+            TaskType.QUICK_ANSWER: 1,
+            TaskType.UNKNOWN_AMBIGUOUS: 2,
+            TaskType.TOOL_HEAVY: 2,
+            TaskType.RETRIEVAL_HEAVY: 2,
+            TaskType.DOCUMENT_DRAFTING: 3,
+            TaskType.CODING: 3,
+            TaskType.RESEARCH: 4,
+            TaskType.LONG_HORIZON: 4,
+            TaskType.LEGAL_REGULATED: 5,
+        }
+        tier = static_map.get(prediction.task_type, 3)
+        models = self.models_by_tier.get(tier, self.models_by_tier[3])
+        model = models[0] if models else list(self.config.models.values())[0]
+        return RoutingDecision(
+            model_id=model.model_id,
+            provider=model.provider,
+            tier=tier,
+            confidence=0.6,
+            reasoning=f"Static routing: {prediction.task_type.value} -> tier {tier}",
+            fallback_model_id=self._next_tier_model(tier).model_id if tier < 5 else None,
+        )
+
+    def _route_prompt_only(self, prediction: TaskPrediction) -> RoutingDecision:
+        """Use prompt heuristics to decide model tier."""
+        # Heuristic: if risk > 0.7 or expected tier >= 4, use frontier
+        if prediction.risk_of_failure > 0.7 or prediction.expected_model_tier >= 4:
+            tier = 4
+        elif prediction.expected_model_tier <= 2:
+            tier = max(prediction.expected_model_tier, 1)
+        else:
+            # Start cheap, escalate on low confidence
+            tier = max(prediction.expected_model_tier - 1, 1)
+        
+        models = self.models_by_tier.get(tier, self.models_by_tier[3])
+        model = models[0] if models else list(self.config.models.values())[0]
+        
+        fallback = None
+        if tier < 5 and prediction.risk_of_failure > 0.5:
+            fallback = self._next_tier_model(tier)
+            fallback_id = fallback.model_id if fallback else None
+        else:
+            fallback_id = None
+        
+        return RoutingDecision(
+            model_id=model.model_id,
+            provider=model.provider,
+            tier=tier,
+            confidence=1.0 - prediction.risk_of_failure,
+            reasoning=f"Prompt heuristic: risk={prediction.risk_of_failure:.2f}, expected_tier={prediction.expected_model_tier}",
+            fallback_model_id=fallback_id,
+            use_verifier=prediction.verifier_required,
+        )
+
+    def _route_learned(self, prediction: TaskPrediction, verifier_fallback: bool = False) -> RoutingDecision:
+        """Learned router with cost-quality tradeoff.
+        
+        In a full implementation, this would load a trained classifier.
+        Here we use a heuristic calibrated from routing_stats.
+        """
+        # Check historical success rate per tier for this task type
+        task_key = prediction.task_type.value
+        best_tier = None
+        best_score = -float("inf")
+        
+        for tier in self.TIER_ORDER:
+            stats = self.routing_stats.get(f"{task_key}_tier_{tier}", {})
+            success_rate = stats.get("success_rate", 0.5)
+            avg_cost = stats.get("avg_cost", 0.01 * tier)
+            
+            # Score = success_weight * success_rate - cost_weight * cost
+            score = 10 * success_rate - 100 * avg_cost
+            
+            # Penalize tiers below expected if risk is high
+            if tier < prediction.expected_model_tier and prediction.risk_of_failure > 0.5:
+                score -= 5
+            
+            if score > best_score:
+                best_score = score
+                best_tier = tier
+        
+        # Default to expected tier if no history
+        if best_tier is None:
+            best_tier = prediction.expected_model_tier
+        
+        models = self.models_by_tier.get(best_tier, self.models_by_tier[3])
+        model = models[0] if models else list(self.config.models.values())[0]
+        
+        # Verifier fallback on uncertain predictions
+        use_verifier = verifier_fallback and prediction.risk_of_failure > 0.5
+        
+        return RoutingDecision(
+            model_id=model.model_id,
+            provider=model.provider,
+            tier=best_tier,
+            confidence=min(best_score / 10 + 0.5, 1.0),
+            reasoning=f"Learned router: tier {best_tier} scored {best_score:.3f} for {task_key}",
+            fallback_model_id=self._next_tier_model(best_tier).model_id if best_tier < 5 else None,
+            use_verifier=use_verifier,
+        )
+
+    def _route_cascade(self, prediction: TaskPrediction) -> RoutingDecision:
+        """FrugalGPT-style cascade: try cheap first, escalate on low confidence."""
+        start_tier = max(1, prediction.expected_model_tier - 2)
+        
+        # Don't start below tier 2 for risky tasks
+        if prediction.risk_of_failure > 0.6:
+            start_tier = max(start_tier, 2)
+        
+        models = self.models_by_tier.get(start_tier, [])
+        if not models:
+            models = self.models_by_tier.get(1, [])
+        if not models:
+            models = list(self.config.models.values())
+        
+        model = models[0]
+        
+        # Determine if we should pre-escalate (for critical tasks)
+        pre_escalate = prediction.task_type == TaskType.LEGAL_REGULATED
+        
+        fallback = None
+        if not pre_escalate and start_tier < prediction.expected_model_tier:
+            fallback = self._next_tier_model(start_tier)
+        
+        return RoutingDecision(
+            model_id=model.model_id,
+            provider=model.provider,
+            tier=start_tier,
+            confidence=1.0 - prediction.risk_of_failure,
+            reasoning=f"Cascade start at tier {start_tier}, expected tier {prediction.expected_model_tier}, risk={prediction.risk_of_failure:.2f}",
+            fallback_model_id=fallback.model_id if fallback else None,
+            use_verifier=prediction.verifier_required,
+        )
+
+    def _next_tier_model(self, current_tier: int) -> Optional[ModelConfig]:
+        for tier in range(current_tier + 1, 6):
+            models = self.models_by_tier.get(tier)
+            if models:
+                return models[0]
+        return None
+
+    def update_stats(self, task_type: TaskType, tier: int, cost: float, success: bool) -> None:
+        key = f"{task_type.value}_tier_{tier}"
+        stats = self.routing_stats.setdefault(key, {"count": 0, "successes": 0, "total_cost": 0.0})
+        stats["count"] += 1
+        if success:
+            stats["successes"] += 1
+        stats["total_cost"] += cost
+        stats["success_rate"] = stats["successes"] / stats["count"]
+        stats["avg_cost"] = stats["total_cost"] / stats["count"]
+
+    def should_escalate(self, decision: RoutingDecision, step_outcome: Outcome, confidence: float) -> bool:
+        """Decide whether to escalate to a stronger model after a step."""
+        if decision.tier >= 5:
+            return False
+        if step_outcome == Outcome.FAILURE and confidence < 0.5:
+            return True
+        if step_outcome == Outcome.PARTIAL_SUCCESS and decision.tier < 4:
+            return True
+        return False