mur_core/

cost.rs

//! Cost tracking and estimation for AI model API calls.
//!
//! Tracks actual spending and provides pre-execution cost estimates
//! so users know what a workflow run will cost before committing.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use crate::model::ModelRole;

/// Cost per million tokens for a model (input and output).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelPricing {
    /// Cost per million input tokens (USD).
    pub input_per_million: f64,
    /// Cost per million output tokens (USD).
    pub output_per_million: f64,
}

/// A single recorded cost event.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostEvent {
    pub model_role: ModelRole,
    pub model_name: String,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cost: f64,
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

/// Tracks API costs across model invocations.
#[derive(Debug, Clone)]
pub struct CostTracker {
    /// Known pricing for models.
    pricing: HashMap<String, ModelPricing>,
    /// Recorded cost events (shared across threads).
    events: Arc<Mutex<Vec<CostEvent>>>,
}

impl CostTracker {
    /// Create a new cost tracker with default model pricing.
    pub fn new() -> Self {
        let mut pricing = HashMap::new();

        // Anthropic models
        pricing.insert(
            "claude-opus-4-5".into(),
            ModelPricing {
                input_per_million: 15.0,
                output_per_million: 75.0,
            },
        );
        pricing.insert(
            "claude-sonnet-4".into(),
            ModelPricing {
                input_per_million: 3.0,
                output_per_million: 15.0,
            },
        );

        // OpenRouter / Google
        pricing.insert(
            "gemini-2.5-flash".into(),
            ModelPricing {
                input_per_million: 0.15,
                output_per_million: 0.60,
            },
        );

        // Local models (free)
        pricing.insert(
            "ollama-local".into(),
            ModelPricing {
                input_per_million: 0.0,
                output_per_million: 0.0,
            },
        );

        Self {
            pricing,
            events: Arc::new(Mutex::new(Vec::new())),
        }
    }

    /// Estimate the cost for a given model and token count.
    pub fn estimate_cost(
        &self,
        model_name: &str,
        input_tokens: u64,
        output_tokens: u64,
    ) -> f64 {
        if let Some(pricing) = self.pricing.get(model_name) {
            let input_cost = (input_tokens as f64 / 1_000_000.0) * pricing.input_per_million;
            let output_cost = (output_tokens as f64 / 1_000_000.0) * pricing.output_per_million;
            input_cost + output_cost
        } else {
            // Unknown model — return a conservative estimate
            let input_cost = (input_tokens as f64 / 1_000_000.0) * 3.0;
            let output_cost = (output_tokens as f64 / 1_000_000.0) * 15.0;
            input_cost + output_cost
        }
    }

    /// Record an actual cost event.
    pub fn record(
        &self,
        model_role: ModelRole,
        model_name: &str,
        input_tokens: u64,
        output_tokens: u64,
    ) {
        let cost = self.estimate_cost(model_name, input_tokens, output_tokens);
        let event = CostEvent {
            model_role,
            model_name: model_name.to_string(),
            input_tokens,
            output_tokens,
            cost,
            timestamp: chrono::Utc::now(),
        };
        if let Ok(mut events) = self.events.lock() {
            events.push(event);
        }
    }

    /// Get total cost across all recorded events.
    pub fn total_cost(&self) -> f64 {
        self.events
            .lock()
            .map(|events| events.iter().map(|e| e.cost).sum())
            .unwrap_or(0.0)
    }

    /// Get cost breakdown by model role.
    pub fn cost_by_role(&self) -> HashMap<ModelRole, f64> {
        let mut breakdown = HashMap::new();
        if let Ok(events) = self.events.lock() {
            for event in events.iter() {
                *breakdown.entry(event.model_role.clone()).or_insert(0.0) += event.cost;
            }
        }
        breakdown
    }

    /// Get all recorded events.
    pub fn events(&self) -> Vec<CostEvent> {
        self.events
            .lock()
            .map(|events| events.clone())
            .unwrap_or_default()
    }

    /// Reset the tracker (clear all events).
    pub fn reset(&self) {
        if let Ok(mut events) = self.events.lock() {
            events.clear();
        }
    }
}

impl Default for CostTracker {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_estimate_cost_known_model() {
        let tracker = CostTracker::new();
        // claude-opus-4-5: $15/M input, $75/M output
        let cost = tracker.estimate_cost("claude-opus-4-5", 1_000_000, 1_000_000);
        assert!((cost - 90.0).abs() < 0.01); // $15 + $75 = $90
    }

    #[test]
    fn test_estimate_cost_cheap_model() {
        let tracker = CostTracker::new();
        // gemini-2.5-flash: $0.15/M input, $0.60/M output
        let cost = tracker.estimate_cost("gemini-2.5-flash", 1_000_000, 1_000_000);
        assert!((cost - 0.75).abs() < 0.01); // $0.15 + $0.60 = $0.75
    }

    #[test]
    fn test_estimate_cost_local_model() {
        let tracker = CostTracker::new();
        let cost = tracker.estimate_cost("ollama-local", 1_000_000, 1_000_000);
        assert!((cost - 0.0).abs() < 0.01);
    }

    #[test]
    fn test_estimate_cost_unknown_model() {
        let tracker = CostTracker::new();
        let cost = tracker.estimate_cost("unknown-model", 1_000_000, 1_000_000);
        // Should use conservative default: $3/M input + $15/M output
        assert!((cost - 18.0).abs() < 0.01);
    }

    #[test]
    fn test_record_and_total() {
        let tracker = CostTracker::new();
        tracker.record(ModelRole::Thinking, "claude-opus-4-5", 500_000, 100_000);
        tracker.record(ModelRole::Task, "gemini-2.5-flash", 200_000, 50_000);

        let total = tracker.total_cost();
        assert!(total > 0.0);

        let events = tracker.events();
        assert_eq!(events.len(), 2);
    }

    #[test]
    fn test_cost_by_role() {
        let tracker = CostTracker::new();
        tracker.record(ModelRole::Thinking, "claude-opus-4-5", 1_000_000, 0);
        tracker.record(ModelRole::Task, "gemini-2.5-flash", 1_000_000, 0);

        let breakdown = tracker.cost_by_role();
        assert!(breakdown.contains_key(&ModelRole::Thinking));
        assert!(breakdown.contains_key(&ModelRole::Task));
        assert!(breakdown[&ModelRole::Thinking] > breakdown[&ModelRole::Task]);
    }

    #[test]
    fn test_reset() {
        let tracker = CostTracker::new();
        tracker.record(ModelRole::Thinking, "claude-opus-4-5", 1_000_000, 0);
        assert!(tracker.total_cost() > 0.0);

        tracker.reset();
        assert!((tracker.total_cost() - 0.0).abs() < 0.001);
        assert!(tracker.events().is_empty());
    }
}