train_router_gen.py

# ─── Synthetic Trace Generator (for training data) ────────────────
MODEL_CONFIGS = {
    "tiny_local": {"tier":1,"cost_input":0.0001,"cost_output":0.0002,"latency":200,"strength":0.35},
    "cheap_cloud": {"tier":2,"cost_input":0.00015,"cost_output":0.0006,"latency":400,"strength":0.55},
    "medium": {"tier":3,"cost_input":0.0015,"cost_output":0.006,"latency":800,"strength":0.80},
    "frontier": {"tier":4,"cost_input":0.005,"cost_output":0.015,"latency":1500,"strength":0.93},
    "specialist": {"tier":5,"cost_input":0.01,"cost_output":0.03,"latency":2000,"strength":0.97},
}
TIER_TO_MODEL = {1:"tiny_local",2:"cheap_cloud",3:"medium",4:"frontier",5:"specialist"}
TIER_COST_MULT = {1:0.05,2:0.15,3:0.75,4:1.0,5:1.5}

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?","When was the Declaration of Independence signed?",
    ],
    "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.",
        "Write a recursive descent parser for JSON.",
    ],
    "research": [
        "Research latest transformer advances.",
        "Find sources comparing LoRA and full FT.",
        "Investigate data center climate impact.",
        "What does literature say on speculative decoding?",
        "Survey privacy-preserving ML techniques.",
        "Compare reinforcement learning algorithms for robotics.",
        "Analyze trends in LLM scaling laws.",
    ],
    "document_drafting": [
        "Draft project proposal for ML pipeline.",
        "Write email to team about deployment.",
        "Create technical report on performance.",
        "Write a project brief for the migration.",
        "Draft meeting notes summary.",
    ],
    "legal_regulated": [
        "Review this contract for liability clauses.",
        "Check GDPR compliance for data pipeline.",
        "Draft privacy policy section.",
        "Analyze indemnification clause in vendor agreement.",
        "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.",
        "Aggregate metrics from 5 monitoring endpoints.",
    ],
    "retrieval_heavy": [
        "Answer based on 50-page document.",
        "Find all payment processing mentions.",
        "Retrieve relevant cases for legal query.",
        "Summarize the quarterly earnings report.",
    ],
    "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?","There's a problem with the data.",
    ],
}

def tier_success_prob(tier, difficulty):
    strength = {1:0.35,2:0.55,3:0.80,4:0.93,5:0.97}.get(tier,0.5)
    return strength ** (difficulty * 0.6)

def generate_training_trace(idx, rng):
    task_types = list(TASK_TEMPLATES.keys())
    task_type = rng.choice(task_types)
    difficulty = {
        "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,
    }[task_type]

    # Try ALL tiers for this task to get ground truth
    tier_outcomes = {}
    for tier in range(1, 6):
        sp = tier_success_prob(tier, difficulty)
        tier_outcomes[tier] = rng.random() < sp

    optimal_tier = 5  # default: need strongest
    for tier in range(1, 6):
        if tier_outcomes.get(tier, False):
            optimal_tier = tier
            break

    actual_tier = rng.choice(list(range(1, 6)))
    # Bias toward reasonable tiers
    if difficulty <= 2:
        actual_tier = rng.choices([1,2,3,4,5], weights=[3,4,2,1,0.5])[0]
    elif difficulty == 3:
        actual_tier = rng.choices([1,2,3,4,5], weights=[1,2,4,2,1])[0]
    elif difficulty == 4:
        actual_tier = rng.choices([1,2,3,4,5], weights=[0.5,1,2,4,2])[0]
    else:
        actual_tier = rng.choices([1,2,3,4,5], weights=[0.2,0.5,1,3,4])[0]

    outcome = "success" if tier_outcomes.get(actual_tier, False) else "failure"

    user_request = rng.choice(TASK_TEMPLATES[task_type])
    cost_mult = TIER_COST_MULT[actual_tier]
    base_tokens = rng.randint(800, 12000) + rng.randint(200, 6000)
    cost = base_tokens / 1000 * MODEL_CONFIGS[TIER_TO_MODEL[actual_tier]]["cost_input"] * cost_mult

    return {
        "trace_id": f"train_{idx}",
        "user_request": user_request,
        "task_type": task_type,
        "difficulty": difficulty,
        "actual_tier": actual_tier,
        "optimal_tier": optimal_tier,
        "outcome": outcome,
        "cost": cost,
        "tier_outcomes": {str(k): v for k, v in tier_outcomes.items()},
        "metadata": {"difficulty": difficulty, "optimal_tier": optimal_tier, "actual_tier": actual_tier},
    }