Vapora/crates/vapora-llm-router/examples/03-cost-tracking.rs

//! # Cost Tracking Example
//!
//! Demonstrates how to track costs across providers and generate reports.
//!
//! ## What This Example Shows
//! - Recording token usage per provider
//! - Calculating costs with provider pricing
//! - Generating cost breakdowns by provider and task type
//! - ROI analysis for optimization
//!
//! ## Run
//! ```bash
//! cargo run --example 03-cost-tracking -p vapora-llm-router
//! ```

fn main() {
    println!("=== Cost Tracking Example ===\n");

    // Step 1: Define provider pricing
    #[derive(Debug, Clone)]
    struct ProviderPricing {
        name: String,
        input_cost_per_1m_tokens: u32,  // cents
        output_cost_per_1m_tokens: u32, // cents
    }

    let providers = vec![
        ProviderPricing {
            name: "claude".to_string(),
            input_cost_per_1m_tokens: 300,   // $3 per 1M
            output_cost_per_1m_tokens: 1500, // $15 per 1M
        },
        ProviderPricing {
            name: "gpt-4".to_string(),
            input_cost_per_1m_tokens: 1000,  // $10 per 1M
            output_cost_per_1m_tokens: 3000, // $30 per 1M
        },
        ProviderPricing {
            name: "gemini".to_string(),
            input_cost_per_1m_tokens: 500,   // $5 per 1M
            output_cost_per_1m_tokens: 1500, // $15 per 1M
        },
    ];

    println!("Provider Pricing (per 1M tokens):\n");
    for p in &providers {
        println!(
            "  {}: input=${:.2}, output=${:.2}",
            p.name,
            p.input_cost_per_1m_tokens as f64 / 100.0,
            p.output_cost_per_1m_tokens as f64 / 100.0
        );
    }

    // Step 2: Record token usage
    println!("\n=== Token Usage by Task ===\n");

    #[derive(Debug, Clone)]
    struct TaskExecution {
        task_id: String,
        provider: String,
        task_type: String,
        input_tokens: u64,
        output_tokens: u64,
    }

    let executions = vec![
        TaskExecution {
            task_id: "task-001".to_string(),
            provider: "claude".to_string(),
            task_type: "coding".to_string(),
            input_tokens: 1500,
            output_tokens: 800,
        },
        TaskExecution {
            task_id: "task-002".to_string(),
            provider: "claude".to_string(),
            task_type: "coding".to_string(),
            input_tokens: 2000,
            output_tokens: 1200,
        },
        TaskExecution {
            task_id: "task-003".to_string(),
            provider: "gpt-4".to_string(),
            task_type: "analysis".to_string(),
            input_tokens: 3000,
            output_tokens: 1500,
        },
        TaskExecution {
            task_id: "task-004".to_string(),
            provider: "gemini".to_string(),
            task_type: "documentation".to_string(),
            input_tokens: 2500,
            output_tokens: 2000,
        },
    ];

    let mut total_cost_cents = 0;

    for exec in &executions {
        let provider = providers.iter().find(|p| p.name == exec.provider).unwrap();
        let input_cost = (exec.input_tokens * provider.input_cost_per_1m_tokens as u64) / 1_000_000;
        let output_cost =
            (exec.output_tokens * provider.output_cost_per_1m_tokens as u64) / 1_000_000;
        let task_cost = input_cost + output_cost;
        total_cost_cents += task_cost;

        println!(
            "{}: {} tokens",
            exec.task_id,
            exec.input_tokens + exec.output_tokens
        );
        println!(
            "  Provider: {}, Task type: {}",
            exec.provider, exec.task_type
        );
        println!(
            "  Cost: ${:.4} (input: ${:.4}, output: ${:.4})\n",
            task_cost as f64 / 10000.0,
            input_cost as f64 / 10000.0,
            output_cost as f64 / 10000.0
        );
    }

    // Step 3: Cost breakdown by provider
    println!("=== Cost Breakdown by Provider ===\n");

    let mut provider_costs: std::collections::HashMap<String, u64> =
        std::collections::HashMap::new();
    let mut provider_tokens: std::collections::HashMap<String, u64> =
        std::collections::HashMap::new();

    for exec in &executions {
        let provider = providers.iter().find(|p| p.name == exec.provider).unwrap();
        let input_cost = (exec.input_tokens * provider.input_cost_per_1m_tokens as u64) / 1_000_000;
        let output_cost =
            (exec.output_tokens * provider.output_cost_per_1m_tokens as u64) / 1_000_000;
        let task_cost = input_cost + output_cost;

        *provider_costs.entry(exec.provider.clone()).or_insert(0) += task_cost;
        *provider_tokens.entry(exec.provider.clone()).or_insert(0) +=
            exec.input_tokens + exec.output_tokens;
    }

    for (provider, cost) in provider_costs.iter() {
        let tokens = provider_tokens.get(provider).unwrap_or(&0);
        let cost_per_1m = if *tokens > 0 {
            (*cost * 1_000_000) / tokens
        } else {
            0
        };
        println!(
            "  {}: ${:.4} ({} tokens, ${:.2}/1M tokens)",
            provider,
            *cost as f64 / 10000.0,
            tokens,
            cost_per_1m as f64 / 100.0
        );
    }

    // Step 4: Cost breakdown by task type
    println!("\n=== Cost Breakdown by Task Type ===\n");

    let mut task_type_costs: std::collections::HashMap<String, u64> =
        std::collections::HashMap::new();

    for exec in &executions {
        let provider = providers.iter().find(|p| p.name == exec.provider).unwrap();
        let input_cost = (exec.input_tokens * provider.input_cost_per_1m_tokens as u64) / 1_000_000;
        let output_cost =
            (exec.output_tokens * provider.output_cost_per_1m_tokens as u64) / 1_000_000;
        let task_cost = input_cost + output_cost;

        *task_type_costs.entry(exec.task_type.clone()).or_insert(0) += task_cost;
    }

    for (task_type, cost) in task_type_costs.iter() {
        let percentage = (*cost as f64 / total_cost_cents as f64) * 100.0;
        println!(
            "  {}: ${:.4} ({:.1}%)",
            task_type,
            *cost as f64 / 10000.0,
            percentage
        );
    }

    // Step 5: Summary
    println!("\n=== Cost Summary ===");
    println!("Total cost: ${:.4}", total_cost_cents as f64 / 10000.0);
    println!(
        "Average cost per task: ${:.4}",
        (total_cost_cents as f64 / executions.len() as f64) / 10000.0
    );
    println!(
        "Total tokens used: {}",
        executions
            .iter()
            .map(|e| e.input_tokens + e.output_tokens)
            .sum::<u64>()
    );
}