training/train_router_v3.py

#!/usr/bin/env python3
"""Trained Router v3: Combined best approach.
Uses per-tier P(success) classifiers + safety floors + 
cost-aware routing with ASYMMETRIC penalties (underkill penalized 5x harder than overkill).
"""
import json, os, sys, random, pickle, uuid
import numpy as np
from datetime import datetime
from collections import defaultdict
from typing import Dict, List, Tuple, Any, Optional

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"]

TIER_STR = {1:0.35,2:0.55,3:0.80,4:0.93,5:0.97}
TIER_COST = {1:0.05,2:0.15,3:0.75,4:1.0,5:1.5}
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}

TASK_TEMPLATES = {
    "quick_answer":["What is the capital of France?","Explain quantum computing briefly.",
        "What is 237*452?","Define photosynthesis.","Who wrote Hamlet?",
        "What is the speed of light?","List the primary colors.","What is GDP?"],
    "coding":["Write a Python function to reverse a linked list.",
        "Fix the bug in this React component.","Refactor auth module to JWT.",
        "Implement LRU cache in Go.","Debug segfault in C++ thread pool.",
        "Add unit tests for the payment module.","Optimize this SQL query.",
        "Create a REST API for user management.","Implement binary search in Rust."],
    "research":["Research latest transformer advances.",
        "Find sources comparing LoRA and full FT.",
        "Investigate data center climate impact.",
        "Survey privacy-preserving ML techniques.",
        "Compare reinforcement learning algorithms for robotics."],
    "document_drafting":["Draft project proposal for ML pipeline.",
        "Write email to team about deployment.","Create technical report on performance."],
    "legal_regulated":["Review this contract for liability clauses.",
        "Check GDPR compliance for data pipeline.","Draft privacy policy section.",
        "Verify regulatory compliance for medical device software."],
    "tool_heavy":["Search open issues and create summary.",
        "Fetch API docs and generate client code.","Query Q3 sales and produce chart."],
    "retrieval_heavy":["Answer based on 50-page document.",
        "Find all payment processing mentions.","Retrieve relevant cases for legal query."],
    "long_horizon":["Plan 3-month roadmap.","Orchestrate multi-region deployment.",
        "Redesign data architecture end-to-end.","Migrate monolith to microservices."],
    "unknown_ambiguous":["Help me with this thing.",
        "I need something about the server.","Can you look into that issue?"],
}

def tsp(tier, diff):
    return TIER_STR[tier] ** (diff * 0.6)

def extract_features(request, task_type, difficulty=3):
    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 gen_trace(idx, rng):
    tt = rng.choice(list(TASK_TEMPLATES.keys()))
    diff = {"quick_answer":1,"document_drafting":2,"tool_heavy":2,"retrieval_heavy":2,
            "research":3,"coding":3,"unknown_ambiguous":3,"long_horizon":4,"legal_regulated":5}[tt]
    tier_out = {}
    for t in range(1,6):
        tier_out[t] = rng.random() < tsp(t, diff)
    opt = 5
    for t in range(1,6):
        if tier_out[t]: opt = t; break
    if diff <= 2: actual = rng.choices([1,2,3,4,5],weights=[3,4,2,1,0.5])[0]
    elif diff == 3: actual = rng.choices([1,2,3,4,5],weights=[1,2,4,2,1])[0]
    elif diff == 4: actual = rng.choices([1,2,3,4,5],weights=[0.5,1,2,4,2])[0]
    else: actual = rng.choices([1,2,3,4,5],weights=[0.2,0.5,1,3,4])[0]
    outcome = "success" if tier_out[actual] else "failure"
    req = rng.choice(TASK_TEMPLATES[tt])
    feats = extract_features(req, tt, diff)
    return {"feats":feats,"opt":opt,"actual":actual,"outcome":outcome,
            "tier_out":tier_out,"tt":tt,"diff":diff,"req":req}

print("="*80)
print("AGENT COST OPTIMIZER - TRAINED ROUTER v3")
print("Asymmetric cost scoring: underkill 5x penalty")
print("="*80)

# ─── Generate ────────────────────────────────────────────────────────
print("\n[1] Generating 50K training traces...")
rng = random.Random(42)
traces = [gen_trace(i, rng) for i in range(50000)]
FEAT_KEYS = sorted(traces[0]["feats"].keys())
def f2v(feats):
    return np.array([float(feats.get(k, 0.0)) for k in FEAT_KEYS], dtype=np.float32)

X_all = np.array([f2v(t["feats"]) for t in traces])
y_opt = np.array([t["opt"] for t in traces])

per_tier_labels = {}
for tier in range(1,6):
    per_tier_labels[tier] = np.array([1 if t["tier_out"].get(tier,False) else 0 for t in traces])

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score
import xgboost as xgb

X_train, X_test, idx_train, idx_test = train_test_split(X_all, range(len(traces)), test_size=0.2, random_state=42, stratify=y_opt)
print(f"  Train: {len(X_train)}, Test: {len(X_test)}")

# ─── Train Per-Tier Classifiers ─────────────────────────────────────
print("\n[2] Training per-tier P(success) classifiers...")
tier_clfs = {}
for tier in range(1,6):
    y_tr = per_tier_labels[tier][idx_train]
    neg = (y_tr==0).sum(); pos = (y_tr==1).sum()
    spw = neg / max(pos,1)
    clf = xgb.XGBClassifier(n_estimators=150, max_depth=5, learning_rate=0.1,
        subsample=0.8, colsample_bytree=0.8, scale_pos_weight=min(spw,5.0),
        objective="binary:logistic", eval_metric="logloss", random_state=42, verbosity=0)
    clf.fit(X_train, y_tr)
    y_pred = clf.predict(X_test)
    acc = accuracy_score(per_tier_labels[tier][idx_test], y_pred)
    f1 = f1_score(per_tier_labels[tier][idx_test], y_pred, zero_division=0)
    tier_clfs[tier] = clf
    print(f"  Tier {tier}: acc={acc:.3f}, f1={f1:.3f}")

# ─── Asymmetric Cost Router ──────────────────────────────────────────
print("\n[3] Building asymmetric cost router...")

def route_asymmetric(x, task_type, tier_clfs, underkill_penalty=5.0, overkill_penalty=1.0):
    """Score each tier with asymmetric penalties.
    
    score(tier) = P(failure@tier) * underkill_penalty * cost_of_failure
                + cost(tier) * overkill_penalty
    
    Underkill (routing too low) is penalized 5x more than overkill.
    """
    if x.ndim == 1:
        x = x.reshape(1, -1)
    floor = TASK_FLOOR.get(task_type, 2)

    best_tier = floor
    best_score = float("inf")

    for tier in range(floor, 6):
        p_fail = 1.0 - tier_clfs[tier].predict_proba(x)[0, 1]
        cost_norm = TIER_COST[tier] / TIER_COST[5]  # [0.03, 1.0]
        
        # Expected cost of failure (cheap model on hard task)
        failure_cost = p_fail * underkill_penalty
        
        # Cost of using this tier (overkill penalty)
        tier_cost = cost_norm * overkill_penalty
        
        score = failure_cost + tier_cost
        
        if score < best_score:
            best_score = score
            best_tier = tier

    return best_tier

# ─── Evaluate ─────────────────────────────────────────────────────────
print("\n[4] Evaluating all routers...")
n_test = len(idx_test)
results = {}

def eval_router(name, route_fn):
    succ = 0; cost = 0.0; unsafe = 0; fd = 0
    td = defaultdict(int)
    for i in idx_test:
        t = traces[i]
        x = f2v(t["feats"]).reshape(1,-1)
        pred = route_fn(x, t)
        td[pred] += 1
        if t["tier_out"].get(pred, False): succ += 1
        elif pred < t["opt"]: unsafe += 1
        elif pred >= t["opt"]: fd += 1
        cost += TIER_COST[pred]
    results[name] = {"success":succ/n_test, "avg_cost":cost/n_test,
                     "unsafe_rate":unsafe/n_test, "false_done":fd/n_test,
                     "tier_dist":dict(td)}

eval_router("always_frontier", lambda x,t: 4)
eval_router("always_cheap", lambda x,t: 1)
eval_router("heuristic_diff+1", lambda x,t: min(t["diff"]+1,5))
eval_router("heuristic_floor", lambda x,t: TASK_FLOOR.get(t["tt"],2))

for ukp in [3.0, 5.0, 8.0, 10.0, 15.0]:
    eval_router(f"asymmetric_uk{ukp:.0f}", lambda x,t,uk=ukp: route_asymmetric(x, t["tt"], tier_clfs, underkill_penalty=uk))

eval_router("oracle", lambda x,t: t["opt"])

# Print comparison
print(f"\n{'Router':<25} {'Success':>10} {'AvgCost':>10} {'CostRed':>10} {'Unsafe':>10} {'F-DONE':>10}")
print("-"*75)
fc = results["always_frontier"]["avg_cost"]
for name, r in sorted(results.items(), key=lambda x: (-x[1]["success"], x[1]["avg_cost"])):
    cr = (1 - r["avg_cost"]/fc)*100
    print(f"{name:<25} {r['success']:>10.3f} {r['avg_cost']:>10.4f} {cr:>9.1f}% {r['unsafe_rate']:>10.3f} {r['false_done']:>10.3f}")

# Pareto frontier
print("\nPARETO FRONTIER:")
pareto = []
for name, r in results.items():
    if name == "always_cheap": continue
    dominated = False
    for name2, r2 in results.items():
        if name == name2: continue
        if r2["success"] >= r["success"] and r2["avg_cost"] <= r["avg_cost"]:
            if r2["success"] > r["success"] or r2["avg_cost"] < r["avg_cost"]:
                dominated = True; break
    if not dominated:
        pareto.append((name, r))
        cr = (1 - r["avg_cost"]/fc)*100
        print(f"  {name:<25} success={r['success']:.3f} cost={r['avg_cost']:.4f} costRed={cr:.1f}% unsafe={r['unsafe_rate']:.3f}")

# ─── Save Best Model ─────────────────────────────────────────────────
print("\n[5] Saving models...")
os.makedirs("/app/router_models", exist_ok=True)
for tier, clf in tier_clfs.items():
    clf.save_model(f"/app/router_models/tier_{tier}_success.json")
with open("/app/router_models/feat_keys.json","w") as f:
    json.dump(FEAT_KEYS, f)
with open("/app/router_models/tier_config.json","w") as f:
    json.dump({"tier_cost":TIER_COST,"tier_str":TIER_STR,"task_floor":TASK_FLOOR}, f, indent=2)
with open("/app/router_models/eval_results_v3.json","w") as f:
    json.dump(results, f, indent=2, default=str)
print(f"  Saved all models to /app/router_models/")

# Find best config
best_name = ""
best_score = -float("inf")
for name, r in results.items():
    if name in ("oracle","always_cheap"): continue
    # Composite: success*20 - cost*50 - unsafe*100
    score = r["success"]*20 - r["avg_cost"]*50 - r["unsafe_rate"]*100
    if score > best_score:
        best_score = score
        best_name = name
print(f"\n  BEST CONFIG: {best_name} (composite score: {best_score:.2f})")
print(f"\nDONE!")