Vapora/crates/vapora-llm-router/src/cost_ranker.rs

use crate::config::ProviderConfig;
use serde::{Deserialize, Serialize};

/// Provider cost and efficiency score for decision making.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderCostScore {
    /// Provider name (claude, gpt4, gemini, ollama)
    pub provider: String,
    /// Estimated cost in cents for the token count
    pub estimated_cost_cents: u32,
    /// Quality score (0.0-1.0) from KG success rate for task type
    pub quality_score: f64,
    /// Cost efficiency: (quality_score * 100) / (cost_cents + 1)
    /// Prevents division by zero for Ollama (free = $0)
    pub cost_efficiency: f64,
}

/// Service for ranking providers by cost efficiency.
pub struct CostRanker;

impl CostRanker {
    /// Estimate cost in cents for token usage on provider.
    /// Formula: (input_tokens * rate_in + output_tokens * rate_out) / 1M * 100
    /// Costs are stored in dollars, converted to cents for calculation.
    pub fn estimate_cost(config: &ProviderConfig, input_tokens: u64, output_tokens: u64) -> u32 {
        // Convert dollar rates to cents
        let input_cost_cents = config.cost_per_1m_input * 100.0;
        let output_cost_cents = config.cost_per_1m_output * 100.0;

        let input_total = (input_tokens as f64 * input_cost_cents) / 1_000_000.0;
        let output_total = (output_tokens as f64 * output_cost_cents) / 1_000_000.0;

        (input_total + output_total).round() as u32
    }

    /// Get quality score for provider + task type.
    /// In practice, queries KG for success rate. For now, uses provided value.
    pub fn get_quality_score(provider: &str, task_type: &str, _quality_data: Option<f64>) -> f64 {
        // Default quality scores until KG integration provides actual metrics
        match (provider, task_type) {
            ("claude", _) => 0.95, // Highest quality
            ("gpt4", _) => 0.92,   // Very good
            ("gemini", _) => 0.88, // Good
            ("ollama", _) => 0.75, // Decent for local
            (_, _) => 0.5,         // Unknown
        }
    }

    /// Rank providers by cost efficiency.
    /// Formula: efficiency = (quality_score * 100) / (cost_cents + 1)
    /// Higher efficiency = better value for money.
    /// Ordered by efficiency descending (best value first).
    pub fn rank_by_efficiency(
        providers: Vec<(String, ProviderConfig)>,
        task_type: &str,
        input_tokens: u64,
        output_tokens: u64,
    ) -> Vec<ProviderCostScore> {
        let mut scores: Vec<ProviderCostScore> = providers
            .into_iter()
            .map(|(provider_name, config)| {
                let cost = Self::estimate_cost(&config, input_tokens, output_tokens);
                let quality = Self::get_quality_score(&provider_name, task_type, None);
                let efficiency = (quality * 100.0) / (cost as f64 + 1.0);

                ProviderCostScore {
                    provider: provider_name,
                    estimated_cost_cents: cost,
                    quality_score: quality,
                    cost_efficiency: efficiency,
                }
            })
            .collect();

        // Sort by efficiency descending (best value first)
        scores.sort_by(|a, b| {
            b.cost_efficiency
                .partial_cmp(&a.cost_efficiency)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        scores
    }

    /// Select cheapest provider when budget is tight.
    /// Orders by cost ascending (cheapest first).
    pub fn rank_by_cost(
        providers: Vec<(String, ProviderConfig)>,
        input_tokens: u64,
        output_tokens: u64,
    ) -> Vec<ProviderCostScore> {
        let mut scores: Vec<ProviderCostScore> = providers
            .into_iter()
            .map(|(provider_name, config)| {
                let cost = Self::estimate_cost(&config, input_tokens, output_tokens);
                let quality = Self::get_quality_score(&provider_name, "generic", None);
                let efficiency = (quality * 100.0) / (cost as f64 + 1.0);

                ProviderCostScore {
                    provider: provider_name,
                    estimated_cost_cents: cost,
                    quality_score: quality,
                    cost_efficiency: efficiency,
                }
            })
            .collect();

        // Sort by cost ascending (cheapest first)
        scores.sort_by_key(|s| s.estimated_cost_cents);

        scores
    }

    /// Calculate cost-benefit ratio for task.
    /// Returns tuple: (provider, cost_cents, efficiency_score)
    pub fn cost_benefit_ratio(
        providers: Vec<(String, ProviderConfig)>,
        task_type: &str,
        input_tokens: u64,
        output_tokens: u64,
    ) -> Vec<(String, u32, f64)> {
        let ranked = Self::rank_by_efficiency(providers, task_type, input_tokens, output_tokens);
        ranked
            .into_iter()
            .map(|score| {
                (
                    score.provider,
                    score.estimated_cost_cents,
                    score.cost_efficiency,
                )
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_quality_score() {
        let claude_quality = CostRanker::get_quality_score("claude", "coding", None);
        assert_eq!(claude_quality, 0.95);

        let ollama_quality = CostRanker::get_quality_score("ollama", "coding", None);
        assert_eq!(ollama_quality, 0.75);

        let unknown_quality = CostRanker::get_quality_score("unknown", "coding", None);
        assert_eq!(unknown_quality, 0.5);
    }
}