cortex_rs_stats/
lib.rs

1// relay-stats: usage + cost dashboard
2// Purpose: fetch rolling 7-day token/cost data from provider APIs or local log scraper
3// Public surface: StatsEngine (new, fetch_summary)
4// NOT responsible for: display/formatting (handled by relay-cli), memory, routing
5// Test strategy: unit test log-scraper against fixture JSONL; API path requires integration test;
6//                README code example is compile-tested via doctest below.
7
8#![doc = include_str!("../README.md")]
9
10mod log_scraper;
11mod provider_api;
12mod savings;
13
14pub use log_scraper::{scrape_claude_logs, TokenCounts};
15
16use anyhow::Result;
17use cortex_rs_core::{Db, CortexConfig};
18use serde::{Deserialize, Serialize};
19use std::collections::HashMap;
20
21pub struct StatsEngine {
22    config: CortexConfig,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct UsageSummary {
27    pub days: u32,
28    pub by_model: Vec<ModelUsage>,
29    pub total_cost_usd: f64,
30    pub routing_savings_usd: f64,
31    pub source: StatsSource,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct ModelUsage {
36    pub model: String,
37    pub input_tokens: u64,
38    pub output_tokens: u64,
39    #[serde(default)]
40    pub cache_creation_input_tokens: u64,
41    #[serde(default)]
42    pub cache_read_input_tokens: u64,
43    pub cost_usd: f64,
44    pub pct: f64,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub enum StatsSource {
49    /// Fetched from provider org API with admin key
50    ProviderApi,
51    /// Parsed from ~/.claude/ session JSONL files
52    LocalLogScraper,
53    /// No data available
54    Unavailable,
55}
56
57impl StatsEngine {
58    pub fn new(config: CortexConfig) -> Self {
59        StatsEngine { config }
60    }
61
62    pub async fn fetch_summary(&self, days: u32) -> Result<UsageSummary> {
63        // Try provider APIs first if admin keys are configured.
64        // Merge Anthropic + OpenAI summaries when both are present.
65        let mut api_models: Vec<ModelUsage> = Vec::new();
66        let mut used_api = false;
67
68        if let Some(key) = self.admin_key("anthropic") {
69            match provider_api::fetch_anthropic(&key, days).await {
70                Ok(s) => {
71                    api_models.extend(s.by_model);
72                    used_api = true;
73                }
74                Err(e) => tracing::warn!("Anthropic API failed: {}", e),
75            }
76        }
77        if let Some(key) = self.admin_key("openai") {
78            match provider_api::fetch_openai(&key, days).await {
79                Ok(s) => {
80                    api_models.extend(s.by_model);
81                    used_api = true;
82                }
83                Err(e) => tracing::warn!("OpenAI API failed: {}", e),
84            }
85        }
86
87        if used_api {
88            // Apply pricing config to any rows the provider API didn't price itself
89            // (OpenAI returns tokens but not cost on the usage endpoint).
90            for m in &mut api_models {
91                if m.cost_usd == 0.0 && (m.input_tokens > 0 || m.output_tokens > 0) {
92                    m.cost_usd = self.model_cost(
93                        &m.model,
94                        TokenCounts {
95                            input: m.input_tokens,
96                            output: m.output_tokens,
97                            cache_creation_input: m.cache_creation_input_tokens,
98                            cache_read_input: m.cache_read_input_tokens,
99                        },
100                    );
101                }
102            }
103            let total = api_models.iter().map(|m| m.cost_usd).sum::<f64>();
104            for m in &mut api_models {
105                m.pct = if total > 0.0 { m.cost_usd / total * 100.0 } else { 0.0 };
106            }
107            api_models.sort_by(|a, b| {
108                b.cost_usd.partial_cmp(&a.cost_usd).unwrap_or(std::cmp::Ordering::Equal)
109            });
110            let routing_savings_usd = self.compute_savings(days).unwrap_or(0.0);
111            return Ok(UsageSummary {
112                days,
113                by_model: api_models,
114                total_cost_usd: total,
115                routing_savings_usd,
116                source: StatsSource::ProviderApi,
117            });
118        }
119
120        // Fallback: parse ~/.claude/ session logs
121        let raw = scrape_claude_logs(days)?;
122        if raw.is_empty() {
123            return Ok(UsageSummary {
124                days,
125                by_model: vec![],
126                total_cost_usd: 0.0,
127                routing_savings_usd: 0.0,
128                source: StatsSource::Unavailable,
129            });
130        }
131
132        let by_model = self.cost_by_model(raw);
133        let total = by_model.iter().map(|m| m.cost_usd).sum::<f64>();
134        let routing_savings_usd = self.compute_savings(days).unwrap_or(0.0);
135
136        Ok(UsageSummary {
137            days,
138            by_model,
139            total_cost_usd: total,
140            routing_savings_usd,
141            source: StatsSource::LocalLogScraper,
142        })
143    }
144
145    fn admin_key(&self, provider: &str) -> Option<String> {
146        let cfg = match provider {
147            "anthropic" => self.config.providers.anthropic.as_ref()?,
148            "openai" => self.config.providers.openai.as_ref()?,
149            _ => return None,
150        };
151        let k = cfg.admin_key.clone()?;
152        if k.is_empty() { None } else { Some(k) }
153    }
154
155    /// Returns estimated USD saved by routing in the last `days`.
156    /// Opens its own DB connection so the engine doesn't need to hold one.
157    fn compute_savings(&self, days: u32) -> Result<f64> {
158        let db_path = self.config.daemon.db_path();
159        if !db_path.exists() {
160            return Ok(0.0);
161        }
162        let db = Db::open(&db_path)?;
163        savings::compute(&db.conn, &self.config, days)
164    }
165
166    fn cost_by_model(&self, raw: HashMap<String, TokenCounts>) -> Vec<ModelUsage> {
167        let total_cost: f64 = raw
168            .iter()
169            .map(|(model, counts)| self.model_cost(model, *counts))
170            .sum();
171
172        let mut usage: Vec<ModelUsage> = raw
173            .into_iter()
174            .map(|(model, counts)| {
175                let cost = self.model_cost(&model, counts);
176                ModelUsage {
177                    model,
178                    input_tokens: counts.input,
179                    output_tokens: counts.output,
180                    cache_creation_input_tokens: counts.cache_creation_input,
181                    cache_read_input_tokens: counts.cache_read_input,
182                    cost_usd: cost,
183                    pct: if total_cost > 0.0 { cost / total_cost * 100.0 } else { 0.0 },
184                }
185            })
186            .collect();
187
188        usage.sort_by(|a, b| {
189            b.cost_usd.partial_cmp(&a.cost_usd).unwrap_or(std::cmp::Ordering::Equal)
190        });
191        usage
192    }
193
194    fn model_cost(&self, model: &str, c: TokenCounts) -> f64 {
195        if let Some(pricing) = self.config.pricing.get(model) {
196            let per = |tokens: u64, rate: f64| (tokens as f64 / 1_000_000.0) * rate;
197            per(c.input, pricing.input)
198                + per(c.output, pricing.output)
199                + per(c.cache_creation_input, pricing.cache_creation_or_default())
200                + per(c.cache_read_input, pricing.cache_read_or_default())
201        } else {
202            // Unknown model: conservative mid-tier estimate (no cache pricing)
203            let per = |tokens: u64, rate: f64| (tokens as f64 / 1_000_000.0) * rate;
204            per(c.input, 3.0) + per(c.output, 15.0) + per(c.cache_read_input, 0.30)
205        }
206    }
207}
cortex_rs_stats/lib.rs

cortex_rs_stats/
lib.rs
