Upload aco/learned_router.py

aco/learned_router.py

@@ -0,0 +1,228 @@
+"""Learned Router - Trainable cost-aware model router.
+
+Trains a simple classifier on synthetic traces to predict optimal model tier.
+Can be swapped into ModelCascadeRouter._route_learned().
+"""
+
+import json
+import pickle
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+from collections import defaultdict
+
+
+@dataclass
+class RouterFeatures:
+    task_type_encoded: int
+    request_length: int
+    num_words: int
+    contains_code_keywords: bool
+    contains_legal_keywords: bool
+    contains_research_keywords: bool
+    contains_tool_keywords: bool
+    expected_difficulty: int
+    has_prior_failures: bool
+    prior_success_rate: float
+
+
+class LearnedRouter:
+    """Simple learned router that can be trained on trace data."""
+
+    def __init__(self):
+        self.weights: Dict[str, float] = {}
+        self.task_type_bias: Dict[str, float] = {}
+        self.trained = False
+
+    def _extract_features(self, user_request: str, task_type: str, history: List[Dict]) -> Dict[str, float]:
+        """Extract features from request and history."""
+        req_lower = user_request.lower()
+        
+        features = {
+            "request_length": len(user_request),
+            "num_words": len(user_request.split()),
+            "has_code": any(kw in req_lower for kw in ["python", "javascript", "code", "function", "bug", "debug", "refactor", "implement", "test"]),
+            "has_legal": any(kw in req_lower for kw in ["contract", "legal", "compliance", "gdpr", "privacy", "policy", "regulatory"]),
+            "has_research": any(kw in req_lower for kw in ["research", "find sources", "literature", "investigate", "compare", "analyze"]),
+            "has_tools": any(kw in req_lower for kw in ["search", "fetch", "retrieve", "query", "api", "database", "scrape"]),
+            "has_long_horizon": any(kw in req_lower for kw in ["plan", "project", "roadmap", "orchestrate", "multi-step"]),
+        }
+        
+        # Historical success rate for this task type
+        similar = [t for t in history if t.get("task_type") == task_type]
+        if similar:
+            success_count = sum(1 for t in similar if t.get("final_outcome") == "success")
+            features["prior_success_rate"] = success_count / len(similar)
+            features["has_prior_failures"] = any(t.get("final_outcome") == "failure" for t in similar[-5:])
+        else:
+            features["prior_success_rate"] = 0.5
+            features["has_prior_failures"] = False
+        
+        return features
+
+    def _score_tier(self, features: Dict[str, float], tier: int) -> float:
+        """Score a tier given features. Higher is better."""
+        if not self.trained:
+            # Heuristic scoring before training
+            base_score = {1: 0.3, 2: 0.5, 3: 0.7, 4: 0.85, 5: 0.9}.get(tier, 0.5)
+            
+            # Adjust by task complexity signals
+            if features["has_legal"] and tier < 4:
+                base_score -= 0.4
+            if features["has_code"] and tier < 3:
+                base_score -= 0.2
+            if features["has_research"] and tier < 3:
+                base_score -= 0.15
+            if features["has_long_horizon"] and tier < 3:
+                base_score -= 0.2
+            
+            # Adjust by history
+            if features["has_prior_failures"] and tier < 3:
+                base_score -= 0.3
+            if features["prior_success_rate"] > 0.8 and tier > 2:
+                base_score += 0.1
+            
+            return base_score
+        
+        # Trained scoring
+        score = self.task_type_bias.get(str(tier), 0.5)
+        for feat_name, feat_val in features.items():
+            weight_key = f"{feat_name}_tier_{tier}"
+            score += self.weights.get(weight_key, 0.0) * (1.0 if feat_val else 0.0)
+        return score
+
+    def predict_tier(self, user_request: str, task_type: str, history: List[Dict] = None) -> Tuple[int, float]:
+        """Predict optimal model tier and confidence."""
+        history = history or []
+        features = self._extract_features(user_request, task_type, history)
+        
+        best_tier = 3
+        best_score = -float("inf")
+        
+        for tier in [1, 2, 3, 4, 5]:
+            score = self._score_tier(features, tier)
+            if score > best_score:
+                best_score = score
+                best_tier = tier
+        
+        confidence = min(best_score, 1.0)
+        return best_tier, confidence
+
+    def train(self, traces: List[Dict]) -> None:
+        """Train router on trace data using simple frequency-based weight learning."""
+        # For each trace, determine optimal tier from outcome
+        tier_counts = defaultdict(lambda: defaultdict(int))
+        
+        for trace in traces:
+            task_type = trace.get("task_type", "unknown")
+            outcome = trace.get("final_outcome", "failure")
+            difficulty = trace.get("metadata", {}).get("difficulty", 3)
+            actual_tier = trace.get("metadata", {}).get("actual_tier", 3)
+            
+            # Optimal tier is the minimum tier that would succeed
+            if outcome == "success":
+                optimal = actual_tier  # This tier succeeded
+            else:
+                optimal = min(actual_tier + 1, 5)  # Need higher tier
+            
+            # Extract features
+            req = trace.get("user_request", "")
+            features = self._extract_features(req, task_type, [])
+            
+            # Count successes per feature+tier combination
+            for feat_name, feat_val in features.items():
+                if isinstance(feat_val, bool) and feat_val:
+                    tier_counts[feat_name][optimal] += 1
+            
+            tier_counts["_overall"][optimal] += 1
+        
+        # Compute bias per tier
+        total = sum(tier_counts["_overall"].values())
+        if total > 0:
+            for tier in [1, 2, 3, 4, 5]:
+                count = tier_counts["_overall"].get(tier, 0)
+                self.task_type_bias[str(tier)] = count / total
+        
+        # Compute feature weights
+        for feat_name, tier_dist in tier_counts.items():
+            if feat_name == "_overall":
+                continue
+            total_feat = sum(tier_dist.values())
+            if total_feat > 0:
+                for tier in [1, 2, 3, 4, 5]:
+                    weight_key = f"{feat_name}_tier_{tier}"
+                    # Positive if this tier is common when feature is present
+                    self.weights[weight_key] = (tier_dist.get(tier, 0) / total_feat) - self.task_type_bias.get(str(tier), 0.1)
+        
+        self.trained = True
+        print(f"Trained router on {len(traces)} traces. Weights: {len(self.weights)}.")
+
+    def save(self, path: str) -> None:
+        with open(path, "wb") as f:
+            pickle.dump({"weights": self.weights, "bias": self.task_type_bias, "trained": self.trained}, f)
+
+    def load(self, path: str) -> None:
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+            self.weights = data["weights"]
+            self.task_type_bias = data["bias"]
+            self.trained = data["trained"]
+
+
+class RouterDatasetBuilder:
+    """Builds training datasets from traces for learned router training."""
+
+    @staticmethod
+    def from_traces(traces: List[Dict]) -> List[Dict]:
+        """Convert traces to (features, optimal_tier) training examples."""
+        examples = []
+        for trace in traces:
+            difficulty = trace.get("metadata", {}).get("difficulty", 3)
+            actual_tier = trace.get("metadata", {}).get("actual_tier", 3)
+            outcome = trace.get("final_outcome", "failure")
+            
+            # Optimal tier
+            if outcome == "success":
+                optimal = actual_tier
+            else:
+                optimal = min(actual_tier + 1, 5)
+            
+            # Simple feature extraction
+            req = trace.get("user_request", "").lower()
+            features = {
+                "length": len(req),
+                "has_code": any(kw in req for kw in ["python", "code", "function", "bug", "debug"]),
+                "has_legal": any(kw in req for kw in ["contract", "legal", "compliance", "gdpr"]),
+                "has_research": any(kw in req for kw in ["research", "find sources", "literature"]),
+                "task_type": trace.get("task_type", "unknown"),
+                "difficulty": difficulty,
+            }
+            
+            examples.append({"features": features, "optimal_tier": optimal, "outcome": outcome})
+        
+        return examples
+
+    @staticmethod
+    def compute_oracle_savings(traces: List[Dict]) -> Dict[str, float]:
+        """Compute what an oracle router (perfect tier selection) would save."""
+        total_cost = 0.0
+        oracle_cost = 0.0
+        
+        for trace in traces:
+            cost = trace.get("total_cost", 0.0)
+            difficulty = trace.get("metadata", {}).get("difficulty", 3)
+            
+            # Oracle uses cheapest tier that can handle difficulty
+            oracle_tier = max(1, difficulty - 1)  # Slightly optimistic
+            oracle_mult = {1: 0.05, 2: 0.15, 3: 0.75, 4: 1.0, 5: 1.5}.get(oracle_tier, 0.75)
+            actual_mult = {1: 0.05, 2: 0.15, 3: 0.75, 4: 1.0, 5: 1.5}.get(
+                trace.get("metadata", {}).get("actual_tier", 3), 0.75)
+            
+            oracle_cost += cost * (oracle_mult / actual_mult) if actual_mult > 0 else cost
+            total_cost += cost
+        
+        return {
+            "total_cost": total_cost,
+            "oracle_cost": oracle_cost,
+            "savings": (total_cost - oracle_cost) / max(total_cost, 1),
+            "per_trace_savings": (total_cost - oracle_cost) / len(traces) if traces else 0,
+        }