Upload aco/learned_router.py with huggingface_hub

aco/learned_router.py

@@ -1,228 +1,151 @@
-"""Learned Router - Trainable cost-aware model router.
+"""Trained Production Router - Replaces heuristic routing.
 
-Trains a simple classifier on synthetic traces to predict optimal model tier.
-Can be swapped into ModelCascadeRouter._route_learned().
+Architecture: difficulty-first + ML confirmation + safety floors.
+
+Usage:
+    from aco.learned_router import TrainedRouter
+    
+    router = TrainedRouter.from_pretrained("narcolepticchicken/agent-cost-optimizer")
+    tier, confidence = router.predict("Write a Python function", "coding", difficulty=3)
 """
 
 import json
+import os
 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()
+try:
+    import numpy as np
+    import xgboost as xgb
+    HAS_ML = True
+except ImportError:
+    HAS_ML = False
+
+
+TASK_TYPES = ["quick_answer","coding","research","document_drafting",
+              "legal_regulated","tool_heavy","retrieval_heavy",
+              "long_horizon","unknown_ambiguous"]
+TT2IDX = {t:i for i,t in enumerate(TASK_TYPES)}
+
+CODE_KW = ["python","javascript","code","function","bug","debug","refactor",
+           "implement","test","compile","runtime","class","module","async","thread"]
+LEGAL_KW = ["contract","legal","compliance","gdpr","privacy","policy","regulatory","liability"]
+RESEARCH_KW = ["research","find sources","literature","investigate","compare","analyze","survey"]
+TOOL_KW = ["search","fetch","retrieve","query","api","database","scrape","aggregate"]
+LONG_KW = ["plan","project","roadmap","orchestrate","multi-step","migrate","pipeline","deploy"]
+MATH_KW = ["calculate","compute","solve","equation","formula","optimize","probability"]
+
+# Default safety floors per task type
+TASK_FLOOR = {
+    "legal_regulated":4,"long_horizon":3,"research":3,"coding":3,
+    "unknown_ambiguous":3,"quick_answer":1,"document_drafting":2,
+    "tool_heavy":2,"retrieval_heavy":2,
+}
+
+
+class TrainedRouter:
+    """Production trained router: difficulty-first + ML confirmation + safety floors."""
+
+    def __init__(self, tier_clfs: Dict, feat_keys: List[str], 
+                 tier_config: Dict, escalation_threshold: float = 0.55):
+        self.tier_clfs = tier_clfs
+        self.feat_keys = feat_keys
+        self.tier_config = tier_config
+        self.tier_cost = {int(k):v for k,v in tier_config["tier_cost"].items()}
+        self.task_floor = tier_config.get("task_floor", TASK_FLOOR)
+        self.escalation_threshold = escalation_threshold
+        self._trained = True
+
+    def extract_features(self, request: str, task_type: str, difficulty: int = 3) -> Dict:
+        r = request.lower()
+        f = {"req_len":len(request),"num_words":len(request.split()),
+            "has_code":int(any(k in r for k in CODE_KW)),
+            "n_code":sum(1 for k in CODE_KW if k in r),
+            "has_legal":int(any(k in r for k in LEGAL_KW)),
+            "n_legal":sum(1 for k in LEGAL_KW if k in r),
+            "has_research":int(any(k in r for k in RESEARCH_KW)),
+            "n_research":sum(1 for k in RESEARCH_KW if k in r),
+            "has_tool":int(any(k in r for k in TOOL_KW)),
+            "n_tool":sum(1 for k in TOOL_KW if k in r),
+            "has_long":int(any(k in r for k in LONG_KW)),
+            "has_math":int(any(k in r for k in MATH_KW)),
+            "tt_idx":TT2IDX.get(task_type,8),"difficulty":difficulty}
+        for tt in TASK_TYPES:
+            f[f"tt_{tt}"] = int(task_type == tt)
+        return f
+
+    def _feats_to_vec(self, feats: Dict):
+        import numpy as np
+        return np.array([float(feats.get(k, 0.0)) for k in self.feat_keys], dtype=np.float32)
+
+    def predict(self, request: str, task_type: str, difficulty: int = 3,
+                escalation_threshold: Optional[float] = None) -> Tuple[int, float]:
+        """Predict optimal tier using difficulty-first + ML confirmation.
+
+        Returns: (tier, confidence)
+        """
+        threshold = escalation_threshold or self.escalation_threshold
         
-        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"]),
-        }
+        # Step 1: difficulty -> base_tier
+        base_tier = min(difficulty + 1, 5)
         
-        # 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
+        # Step 2: apply safety floor
+        floor = self.task_floor.get(task_type, 2)
+        base_tier = max(base_tier, floor)
         
-        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
+        if not HAS_ML or not self._trained:
+            return base_tier, 0.6
 
-    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))
+        # Step 3: ML confirmation
+        feats = self.extract_features(request, task_type, difficulty)
+        x = self._feats_to_vec(feats).reshape(1, -1)
         
-        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
+        p_success = self.tier_clfs[base_tier].predict_proba(x)[0, 1]
+        confidence = p_success
         
-        # 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
+        # Step 4: escalate if P(success) too low
+        while p_success < threshold and base_tier < 5:
+            base_tier += 1
+            p_success = self.tier_clfs[base_tier].predict_proba(x)[0, 1]
+            confidence = p_success
         
-        # 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)
+        return base_tier, float(confidence)
+
+    @classmethod
+    def from_pretrained(cls, repo_id: str, escalation_threshold: float = 0.55,
+                        cache_dir: Optional[str] = None):
+        """Load trained router from HuggingFace Hub."""
+        from huggingface_hub import hf_hub_download
         
-        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})
+        bundle_path = hf_hub_download(
+            repo_id=repo_id, filename="router_models/router_bundle.pkl",
+            cache_dir=cache_dir,
+        )
         
-        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
+        with open(bundle_path, "rb") as f:
+            import pickle
+            bundle = pickle.load(f)
         
-        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 cls(
+            tier_clfs={int(k): v for k, v in bundle["tier_clfs"].items()},
+            feat_keys=bundle["feat_keys"],
+            tier_config=bundle["tier_config"],
+            escalation_threshold=escalation_threshold,
+        )
+
+    @classmethod
+    def from_local(cls, model_dir: str, escalation_threshold: float = 0.55):
+        """Load from local directory."""
+        bundle_path = os.path.join(model_dir, "router_bundle.pkl")
+        with open(bundle_path, "rb") as f:
+            import pickle
+            bundle = pickle.load(f)
         
-        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,
-        }
+        return cls(
+            tier_clfs={int(k): v for k, v in bundle["tier_clfs"].items()},
+            feat_keys=bundle["feat_keys"],
+            tier_config=bundle["tier_config"],
+            escalation_threshold=escalation_threshold,
+        )