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        let total_cost: f64 = costs.iter().map(|c| c.cost_usd).sum();
137        let total_tokens: usize = costs.iter().map(|c| c.input_tokens + c.output_tokens).sum();
138
139        let mut by_model: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
140        for cost in costs.iter() {
141            *by_model.entry(cost.model.clone()).or_insert(0.0) += cost.cost_usd;
142        }
143
144        CostSummary {
145            total_cost,
146            total_tokens,
147            by_model,
148            call_count: costs.len(),
149        }
150    }
151
152    /// Get cost history (with date filtering)
153    pub async fn history(&self, days: usize) -> Vec<CostRecord> {
154        let costs = self.costs.read().await;
155        let cutoff = chrono::Utc::now() - chrono::Duration::days(days as i64);
156
157        costs
158            .iter()
159            .filter(|c| c.timestamp > cutoff)
160            .cloned()
161            .collect()
162    }
163
164    /// Update rate limit status from Anthropic API response headers
165    pub async fn update_rate_limits(&self, headers: &reqwest::header::HeaderMap) {
166        let parse_usize = |key| {
167            headers
168                .get(key)
169                .and_then(|v| v.to_str().ok())
170                .and_then(|s| s.parse().ok())
171        };
172
173        let parse_string = |key| {
174            headers
175                .get(key)
176                .and_then(|v| v.to_str().ok())
177                .map(|s| s.to_string())
178        };
179
180        let status = RateLimitStatus {
181            requests_limit: parse_usize("anthropic-ratelimit-requests-limit"),
182            requests_remaining: parse_usize("anthropic-ratelimit-requests-remaining"),
183            input_tokens_limit: parse_usize("anthropic-ratelimit-input-tokens-limit"),
184            input_tokens_remaining: parse_usize("anthropic-ratelimit-input-tokens-remaining"),
185            output_tokens_limit: parse_usize("anthropic-ratelimit-output-tokens-limit"),
186            output_tokens_remaining: parse_usize("anthropic-ratelimit-output-tokens-remaining"),
187            tokens_reset: parse_string("anthropic-ratelimit-tokens-reset"),
188        };
189
190        *self.rate_limit_status.write().await = Some(status);
191    }
192
193    /// Get current rate limit status
194    pub async fn get_rate_limits(&self) -> Option<RateLimitStatus> {
195        self.rate_limit_status.read().await.clone()
196    }
197}
198
199impl Default for CostTracker {
200    fn default() -> Self {
201        Self::new()
202    }
203}
204
205#[derive(Debug, Serialize, Deserialize)]
206pub struct CostSummary {
207    pub total_cost: f64,
208    pub total_tokens: usize,
209    pub by_model: std::collections::HashMap<String, f64>,
210    pub call_count: usize,
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    #[test]
218    fn test_token_cost_calculation() {
219        let cost = ModelCost {
220            model: "test".to_string(),
221            input_cost_per_1m: 1.0,
222            output_cost_per_1m: 2.0,
223        };
224
225        let usage = TokenUsage {
226            input_tokens: 1_000_000,
227            output_tokens: 1_000_000,
228            total_tokens: 2_000_000,
229        };
230
231        let calculated = usage.calculate_cost(&cost);
232        assert_eq!(calculated, 3.0); // 1.0 + 2.0
233    }
234
235    #[tokio::test]
236    async fn test_cost_tracking() {
237        let tracker = CostTracker::new();
238
239        tracker
240            .record(
241                "claude-opus-4-6",
242                TokenUsage {
243                    input_tokens: 100,
244                    output_tokens: 50,
245                    total_tokens: 150,
246                },
247            )
248            .await
249            .unwrap();
250
251        let summary = tracker.summary().await;
252        assert_eq!(summary.call_count, 1);
253        assert!(summary.total_cost > 0.0);
254    }
255
256    #[tokio::test]
257    async fn test_update_rate_limits() {
258        let tracker = CostTracker::new();
259        let mut headers = reqwest::header::HeaderMap::new();
260
261        headers.insert(
262            "anthropic-ratelimit-requests-limit",
263            "1000".parse().unwrap(),
264        );
265        headers.insert(
266            "anthropic-ratelimit-requests-remaining",
267            "999".parse().unwrap(),
268        );
269        headers.insert(
270            "anthropic-ratelimit-input-tokens-limit",
271            "400000".parse().unwrap(),
272        );
273        headers.insert(
274            "anthropic-ratelimit-input-tokens-remaining",
275            "399000".parse().unwrap(),
276        );
277        headers.insert(
278            "anthropic-ratelimit-output-tokens-limit",
279            "100000".parse().unwrap(),
280        );
281        headers.insert(
282            "anthropic-ratelimit-output-tokens-remaining",
283            "99000".parse().unwrap(),
284        );
285        headers.insert(
286            "anthropic-ratelimit-tokens-reset",
287            "2023-11-20T12:00:00Z".parse().unwrap(),
288        );
289
290        tracker.update_rate_limits(&headers).await;
291
292        let status = tracker.get_rate_limits().await.unwrap();
293
294        assert_eq!(status.requests_limit, Some(1000));
295        assert_eq!(status.requests_remaining, Some(999));
296        assert_eq!(status.input_tokens_limit, Some(400000));
297        assert_eq!(status.input_tokens_remaining, Some(399000));
298        assert_eq!(status.output_tokens_limit, Some(100000));
299        assert_eq!(status.output_tokens_remaining, Some(99000));
300        assert_eq!(
301            status.tokens_reset,
302            Some("2023-11-20T12:00:00Z".to_string())
303        );
304    }
305}
apollo/cost.rs

apollo/
cost.rs
