apollo/
cost.rs

1//! Cost tracking — token counting and billing for LLM calls
2//! Phase 4 feature: Production billing support
3
4use serde::{Deserialize, Serialize};
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8/// Cost per 1M tokens (input/output separate)
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct ModelCost {
11    pub model: String,
12    pub input_cost_per_1m: f64,
13    pub output_cost_per_1m: f64,
14}
15
16/// Token usage for a call
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct TokenUsage {
19    pub input_tokens: usize,
20    pub output_tokens: usize,
21    pub total_tokens: usize,
22}
23
24impl TokenUsage {
25    pub fn calculate_cost(&self, cost: &ModelCost) -> f64 {
26        let input_cost = (self.input_tokens as f64 / 1_000_000.0) * cost.input_cost_per_1m;
27        let output_cost = (self.output_tokens as f64 / 1_000_000.0) * cost.output_cost_per_1m;
28        input_cost + output_cost
29    }
30}
31
32/// Cost record for a single LLM call
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct CostRecord {
35    pub id: String,
36    pub model: String,
37    pub input_tokens: usize,
38    pub output_tokens: usize,
39    pub cost_usd: f64,
40    pub timestamp: chrono::DateTime<chrono::Utc>,
41}
42
43/// Claude API rate limit status (from response headers)
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct RateLimitStatus {
46    pub requests_limit: Option<usize>,
47    pub requests_remaining: Option<usize>,
48    pub input_tokens_limit: Option<usize>,
49    pub input_tokens_remaining: Option<usize>,
50    pub output_tokens_limit: Option<usize>,
51    pub output_tokens_remaining: Option<usize>,
52    pub tokens_reset: Option<String>,
53}
54
55/// Cost tracker (in-memory + persistent accounting hooks)
56pub struct CostTracker {
57    costs: Arc<RwLock<Vec<CostRecord>>>,
58    models: Arc<RwLock<Vec<ModelCost>>>,
59    rate_limit_status: Arc<RwLock<Option<RateLimitStatus>>>,
60}
61
62impl CostTracker {
63    pub fn new() -> Self {
64        let models = vec![
65            ModelCost {
66                model: "claude-opus-4-6".to_string(),
67                input_cost_per_1m: 15.0,
68                output_cost_per_1m: 75.0,
69            },
70            ModelCost {
71                model: "claude-3-5-sonnet-20241022".to_string(),
72                input_cost_per_1m: 3.0,
73                output_cost_per_1m: 15.0,
74            },
75            ModelCost {
76                model: "gpt-4-turbo".to_string(),
77                input_cost_per_1m: 10.0,
78                output_cost_per_1m: 30.0,
79            },
80            ModelCost {
81                model: "gpt-4".to_string(),
82                input_cost_per_1m: 30.0,
83                output_cost_per_1m: 60.0,
84            },
85            ModelCost {
86                model: "gpt-3.5-turbo".to_string(),
87                input_cost_per_1m: 0.5,
88                output_cost_per_1m: 1.5,
89            },
90            ModelCost {
91                model: "gemini-2.0-flash".to_string(),
92                input_cost_per_1m: 0.075,
93                output_cost_per_1m: 0.3,
94            },
95        ];
96
97        Self {
98            costs: Arc::new(RwLock::new(Vec::new())),
99            models: Arc::new(RwLock::new(models)),
100            rate_limit_status: Arc::new(RwLock::new(None)),
101        }
102    }
103
104    /// Record a cost from an LLM call
105    pub async fn record(&self, model: &str, usage: TokenUsage) -> anyhow::Result<()> {
106        let models = self.models.read().await;
107        let model_cost = models
108            .iter()
109            .find(|m| m.model == model)
110            .cloned()
111            .unwrap_or_else(|| ModelCost {
112                model: model.to_string(),
113                input_cost_per_1m: 0.0,
114                output_cost_per_1m: 0.0,
115            });
116
117        let cost_usd = usage.calculate_cost(&model_cost);
118
119        let record = CostRecord {
120            id: uuid::Uuid::new_v4().to_string(),
121            model: model.to_string(),
122            input_tokens: usage.input_tokens,
123            output_tokens: usage.output_tokens,
124            cost_usd,
125            timestamp: chrono::Utc::now(),
126        };
127
128        self.costs.write().await.push(record);
129        Ok(())
130    }
131
132    /// Get cost summary
133    pub async fn summary(&self) -> CostSummary {
134        let costs = self.costs.read().await;
135
136        // Folded from a positive zero rather than summed: `<f64 as Sum>` starts
137        // at -0.0, so an empty tracker would otherwise report "-0.0" spent.
138        let total_cost: f64 = costs.iter().map(|c| c.cost_usd).fold(0.0, |acc, c| acc + c);
139        let total_tokens: usize = costs.iter().map(|c| c.input_tokens + c.output_tokens).sum();
140
141        let mut by_model: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
142        for cost in costs.iter() {
143            *by_model.entry(cost.model.clone()).or_insert(0.0) += cost.cost_usd;
144        }
145
146        CostSummary {
147            total_cost,
148            total_tokens,
149            by_model,
150            call_count: costs.len(),
151        }
152    }
153
154    /// Get cost history (with date filtering)
155    pub async fn history(&self, days: usize) -> Vec<CostRecord> {
156        let costs = self.costs.read().await;
157        let cutoff = chrono::Utc::now() - chrono::Duration::days(days as i64);
158
159        costs
160            .iter()
161            .filter(|c| c.timestamp > cutoff)
162            .cloned()
163            .collect()
164    }
165
166    /// Update rate limit status from Anthropic API response headers
167    pub async fn update_rate_limits(&self, headers: &reqwest::header::HeaderMap) {
168        let parse_usize = |key| {
169            headers
170                .get(key)
171                .and_then(|v| v.to_str().ok())
172                .and_then(|s| s.parse().ok())
173        };
174
175        let parse_string = |key| {
176            headers
177                .get(key)
178                .and_then(|v| v.to_str().ok())
179                .map(|s| s.to_string())
180        };
181
182        let status = RateLimitStatus {
183            requests_limit: parse_usize("anthropic-ratelimit-requests-limit"),
184            requests_remaining: parse_usize("anthropic-ratelimit-requests-remaining"),
185            input_tokens_limit: parse_usize("anthropic-ratelimit-input-tokens-limit"),
186            input_tokens_remaining: parse_usize("anthropic-ratelimit-input-tokens-remaining"),
187            output_tokens_limit: parse_usize("anthropic-ratelimit-output-tokens-limit"),
188            output_tokens_remaining: parse_usize("anthropic-ratelimit-output-tokens-remaining"),
189            tokens_reset: parse_string("anthropic-ratelimit-tokens-reset"),
190        };
191
192        *self.rate_limit_status.write().await = Some(status);
193    }
194
195    /// Get current rate limit status
196    pub async fn get_rate_limits(&self) -> Option<RateLimitStatus> {
197        self.rate_limit_status.read().await.clone()
198    }
199}
200
201impl Default for CostTracker {
202    fn default() -> Self {
203        Self::new()
204    }
205}
206
207#[derive(Debug, Serialize, Deserialize)]
208pub struct CostSummary {
209    pub total_cost: f64,
210    pub total_tokens: usize,
211    pub by_model: std::collections::HashMap<String, f64>,
212    pub call_count: usize,
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    #[test]
220    fn test_token_cost_calculation() {
221        let cost = ModelCost {
222            model: "test".to_string(),
223            input_cost_per_1m: 1.0,
224            output_cost_per_1m: 2.0,
225        };
226
227        let usage = TokenUsage {
228            input_tokens: 1_000_000,
229            output_tokens: 1_000_000,
230            total_tokens: 2_000_000,
231        };
232
233        let calculated = usage.calculate_cost(&cost);
234        assert_eq!(calculated, 3.0); // 1.0 + 2.0
235    }
236
237    #[tokio::test]
238    async fn test_cost_tracking() {
239        let tracker = CostTracker::new();
240
241        tracker
242            .record(
243                "claude-opus-4-6",
244                TokenUsage {
245                    input_tokens: 100,
246                    output_tokens: 50,
247                    total_tokens: 150,
248                },
249            )
250            .await
251            .unwrap();
252
253        let summary = tracker.summary().await;
254        assert_eq!(summary.call_count, 1);
255        assert!(summary.total_cost > 0.0);
256    }
257
258    #[tokio::test]
259    async fn an_empty_tracker_reports_a_positive_zero_cost() {
260        let summary = CostTracker::new().summary().await;
261        assert_eq!(summary.call_count, 0);
262        assert_eq!(summary.total_cost, 0.0);
263        assert!(!summary.total_cost.is_sign_negative());
264        assert_eq!(format!("{:.1}", summary.total_cost), "0.0");
265    }
266
267    #[tokio::test]
268    async fn test_update_rate_limits() {
269        let tracker = CostTracker::new();
270        let mut headers = reqwest::header::HeaderMap::new();
271
272        headers.insert(
273            "anthropic-ratelimit-requests-limit",
274            "1000".parse().unwrap(),
275        );
276        headers.insert(
277            "anthropic-ratelimit-requests-remaining",
278            "999".parse().unwrap(),
279        );
280        headers.insert(
281            "anthropic-ratelimit-input-tokens-limit",
282            "400000".parse().unwrap(),
283        );
284        headers.insert(
285            "anthropic-ratelimit-input-tokens-remaining",
286            "399000".parse().unwrap(),
287        );
288        headers.insert(
289            "anthropic-ratelimit-output-tokens-limit",
290            "100000".parse().unwrap(),
291        );
292        headers.insert(
293            "anthropic-ratelimit-output-tokens-remaining",
294            "99000".parse().unwrap(),
295        );
296        headers.insert(
297            "anthropic-ratelimit-tokens-reset",
298            "2023-11-20T12:00:00Z".parse().unwrap(),
299        );
300
301        tracker.update_rate_limits(&headers).await;
302
303        let status = tracker.get_rate_limits().await.unwrap();
304
305        assert_eq!(status.requests_limit, Some(1000));
306        assert_eq!(status.requests_remaining, Some(999));
307        assert_eq!(status.input_tokens_limit, Some(400000));
308        assert_eq!(status.input_tokens_remaining, Some(399000));
309        assert_eq!(status.output_tokens_limit, Some(100000));
310        assert_eq!(status.output_tokens_remaining, Some(99000));
311        assert_eq!(
312            status.tokens_reset,
313            Some("2023-11-20T12:00:00Z".to_string())
314        );
315    }
316}
apollo/cost.rs

apollo/
cost.rs
