use std::sync::Arc; use prometheus::{GaugeVec, IntCounterVec, Registry}; /// Prometheus metrics for cost tracking and budget enforcement. /// Exposes budget utilization, spending, and fallback events. pub struct CostMetrics { /// Remaining budget per role in cents (gauge) pub budget_remaining_cents: GaugeVec, /// Budget utilization per role (0.0-1.0) (gauge) pub budget_utilization: GaugeVec, /// Cost per provider in cents (counter) pub cost_per_provider_cents: IntCounterVec, /// Fallback triggered events with reason (counter) pub fallback_triggered_total: IntCounterVec, /// Total tokens used per provider (counter) pub tokens_per_provider: IntCounterVec, } impl CostMetrics { /// Create new cost metrics collection (registers with default global /// registry) pub fn new() -> Result, prometheus::Error> { let registry = prometheus::default_registry(); Self::with_registry(registry) } /// Create metrics with existing registry pub fn with_registry(registry: &Registry) -> Result, prometheus::Error> { let budget_remaining_cents = GaugeVec::new( prometheus::Opts::new( "vapora_llm_budget_remaining_cents", "Remaining budget for agent role in cents", ), &["role"], )?; registry.register(Box::new(budget_remaining_cents.clone()))?; let budget_utilization = GaugeVec::new( prometheus::Opts::new( "vapora_llm_budget_utilization", "Budget utilization percentage for agent role (0.0-1.0)", ), &["role"], )?; registry.register(Box::new(budget_utilization.clone()))?; let cost_per_provider_cents = IntCounterVec::new( prometheus::Opts::new( "vapora_llm_cost_per_provider_cents", "Total cost per provider in cents", ), &["provider"], )?; registry.register(Box::new(cost_per_provider_cents.clone()))?; let fallback_triggered_total = IntCounterVec::new( prometheus::Opts::new( "vapora_llm_fallback_triggered_total", "Total times fallback provider was triggered", ), &["role", "reason"], )?; registry.register(Box::new(fallback_triggered_total.clone()))?; let tokens_per_provider = IntCounterVec::new( prometheus::Opts::new( "vapora_llm_tokens_per_provider", "Total tokens processed per provider", ), &["provider", "token_type"], )?; registry.register(Box::new(tokens_per_provider.clone()))?; Ok(Arc::new(Self { budget_remaining_cents, budget_utilization, cost_per_provider_cents, fallback_triggered_total, tokens_per_provider, })) } /// Record budget update for role pub fn record_budget_update(&self, role: &str, remaining_cents: u32, utilization: f64) { self.budget_remaining_cents .with_label_values(&[role]) .set(remaining_cents as f64); self.budget_utilization .with_label_values(&[role]) .set(utilization); } /// Record cost for provider pub fn record_provider_cost(&self, provider: &str, cost_cents: u32) { self.cost_per_provider_cents .with_label_values(&[provider]) .inc_by(cost_cents as u64); } /// Record fallback provider activation pub fn record_fallback_triggered(&self, role: &str, reason: &str) { self.fallback_triggered_total .with_label_values(&[role, reason]) .inc(); } /// Record tokens used per provider pub fn record_tokens(&self, provider: &str, input_tokens: u64, output_tokens: u64) { self.tokens_per_provider .with_label_values(&[provider, "input"]) .inc_by(input_tokens); self.tokens_per_provider .with_label_values(&[provider, "output"]) .inc_by(output_tokens); } } #[cfg(test)] mod tests { use super::*; fn create_test_metrics() -> Arc { let registry = Registry::new(); CostMetrics::with_registry(®istry).expect("Failed to create test metrics") } #[test] fn test_cost_metrics_creation() { let registry = Registry::new(); let metrics = CostMetrics::with_registry(®istry); assert!(metrics.is_ok()); } #[test] fn test_record_budget_update() { let metrics = create_test_metrics(); metrics.record_budget_update("developer", 25000, 0.167); // Metric recorded (would verify via Prometheus gather in integration // test) } #[test] fn test_record_provider_cost() { let metrics = create_test_metrics(); metrics.record_provider_cost("claude", 500); metrics.record_provider_cost("claude", 300); // Counter incremented by 800 total } #[test] fn test_record_fallback_triggered() { let metrics = create_test_metrics(); metrics.record_fallback_triggered("developer", "budget_exceeded"); metrics.record_fallback_triggered("architect", "budget_exceeded"); metrics.record_fallback_triggered("developer", "budget_near_threshold"); // Multiple fallback events recorded } #[test] fn test_record_tokens() { let metrics = create_test_metrics(); metrics.record_tokens("claude", 5000, 1000); metrics.record_tokens("gpt4", 3000, 500); // Token counts recorded per provider } }