Vapora/crates/vapora-llm-router/examples/02-budget-enforcement.rs

//! # Budget Enforcement Example
//!
//! Demonstrates cost control with per-role budget limits and automatic
//! fallback.
//!
//! ## What This Example Shows
//! - Setting monthly and weekly budget limits
//! - Tracking spending against budgets
//! - Three-tier enforcement (normal, near-threshold, exceeded)
//! - Automatic fallback to cheaper providers
//! - Alert thresholds for proactive management
//!
//! ## Run
//! ```bash
//! cargo run --example 02-budget-enforcement -p vapora-llm-router
//! ```

fn main() {
    println!("=== Budget Enforcement Example ===\n");

    // Step 1: Define budget limits for roles
    #[derive(Debug, Clone)]
    struct RoleBudget {
        role: String,
        monthly_limit_cents: u32,
        weekly_limit_cents: u32,
        fallback_provider: String,
    }

    let budgets = vec![
        RoleBudget {
            role: "architect".to_string(),
            monthly_limit_cents: 50000, // $500/month
            weekly_limit_cents: 12500,  // $125/week
            fallback_provider: "gpt-4".to_string(),
        },
        RoleBudget {
            role: "developer".to_string(),
            monthly_limit_cents: 30000, // $300/month
            weekly_limit_cents: 7500,   // $75/week
            fallback_provider: "ollama".to_string(),
        },
    ];

    println!("Budget Configuration:\n");
    for budget in &budgets {
        println!(
            "  {}: ${:.2}/month, ${:.2}/week",
            budget.role,
            budget.monthly_limit_cents as f64 / 100.0,
            budget.weekly_limit_cents as f64 / 100.0
        );
        println!("    Fallback: {}\n", budget.fallback_provider);
    }

    // Step 2: Simulate spending for a developer
    println!("=== Spending Simulation (Developer Role) ===\n");

    let mut monthly_spent_cents = 0;
    let mut transactions = vec![];

    // Simulate task executions with costs
    let task_costs = vec![
        ("analyze_code", 2000, "Claude (Opus)"),
        ("write_docs", 1500, "Claude (Sonnet)"),
        ("run_tests", 500, "Ollama (local)"),
        ("refactor_module", 3000, "Claude (Opus)"),
        ("code_review", 1800, "Claude"),
        ("implement_feature", 4500, "Claude (Opus)"),
        ("fix_bug", 1200, "Claude (Sonnet)"),
    ];

    println!("Task Executions:");
    for (task, cost, provider) in &task_costs {
        monthly_spent_cents += cost;
        transactions.push((task, cost, provider));

        let percentage_used = (monthly_spent_cents as f64 / 30000.0) * 100.0;
        println!(
            "  {} ({}): ${:.2} → Running total: ${:.2} ({:.0}%)",
            task,
            provider,
            *cost as f64 / 100.0,
            monthly_spent_cents as f64 / 100.0,
            percentage_used
        );

        // Determine status
        if percentage_used >= 100.0 {
            println!("    ⚠️  BUDGET EXCEEDED - Switching to fallback provider!");
        } else if percentage_used >= 80.0 {
            println!(
                "    ⚠️  NEAR THRESHOLD ({:.0}%) - Alert triggered",
                percentage_used
            );
        } else if percentage_used >= 50.0 {
            println!("    ℹ️  Halfway through budget ({:.0}%)", percentage_used);
        }
        println!();
    }

    // Step 3: Demonstrate budget enforcement
    println!("=== Budget Status ===");
    let monthly_limit = 30000;
    let percentage = (monthly_spent_cents as f64 / monthly_limit as f64) * 100.0;

    println!("Monthly budget: ${:.2}", monthly_limit as f64 / 100.0);
    println!("Amount spent: ${:.2}", monthly_spent_cents as f64 / 100.0);
    println!(
        "Amount remaining: ${:.2}",
        (monthly_limit - monthly_spent_cents) as f64 / 100.0
    );
    println!("Utilization: {:.1}%\n", percentage);

    // Step 4: Show enforcement levels
    println!("=== Three-Tier Enforcement ===");

    let tiers = vec![
        (0.0, 50.0, "Normal", "Use preferred provider"),
        (50.0, 80.0, "Caution", "Monitor spending closely"),
        (80.0, 100.0, "Near Threshold", "Use cheaper alternative"),
        (100.0, 150.0, "Exceeded", "Use fallback provider only"),
    ];

    for (min, max, tier, action) in tiers {
        let status = if percentage >= min && percentage < max {
            "← CURRENT"
        } else {
            ""
        };
        println!(
            "  {}%–{}%: {} → {}  {}",
            min as u32, max as u32, tier, action, status
        );
    }

    // Step 5: Recommendations
    println!("\n=== Budget Management Strategy ===");
    println!(
        "Current tier: {}",
        if percentage < 50.0 {
            "Normal"
        } else if percentage < 80.0 {
            "Caution"
        } else if percentage < 100.0 {
            "Near Threshold"
        } else {
            "Exceeded"
        }
    );

    if percentage < 80.0 {
        println!("✓ Safe to continue with current spending patterns");
    } else {
        println!("⚠️  Recommendations:");
        println!("  1. Use cheaper models for routine tasks");
        println!("  2. Use Ollama (local) for testing/feedback");
        println!("  3. Batch similar tasks for efficiency");
        println!("  4. Consider increasing budget for next month");
    }
}