usage_monitor_cli/provider/
openai.rs

1use async_trait::async_trait;
2use chrono::Utc;
3
4use crate::error::SpendPanelError;
5use crate::model::{CostSnapshot, DailyCost, PlanInfo, RateWindow, SpendLimit, UsageSnapshot};
6use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
7
8/// Rate limit headers returned by OpenAI.
9#[derive(Debug, Default)]
10struct OpenAIRateLimitHeaders {
11    limit_requests: Option<u64>,
12    remaining_requests: Option<u64>,
13    limit_tokens: Option<u64>,
14    remaining_tokens: Option<u64>,
15}
16
17impl OpenAIRateLimitHeaders {
18    fn from_headers(headers: &reqwest::header::HeaderMap) -> Self {
19        let parse_u64 = |name: &str| -> Option<u64> {
20            headers
21                .get(name)
22                .and_then(|v| v.to_str().ok())
23                .and_then(|s| s.parse().ok())
24        };
25
26        Self {
27            limit_requests: parse_u64("x-ratelimit-limit-requests"),
28            remaining_requests: parse_u64("x-ratelimit-remaining-requests"),
29            limit_tokens: parse_u64("x-ratelimit-limit-tokens"),
30            remaining_tokens: parse_u64("x-ratelimit-remaining-tokens"),
31        }
32    }
33}
34
35/// OpenAI cost API response structure.
36#[derive(serde::Deserialize, Debug)]
37struct OpenAICostResponse {
38    data: Vec<OpenAICostItem>,
39}
40
41#[derive(serde::Deserialize, Debug)]
42struct OpenAICostItem {
43    amount: OpenAICostAmount,
44    #[allow(dead_code)]
45    line_item: Option<String>,
46}
47
48#[derive(serde::Deserialize, Debug)]
49struct OpenAICostAmount {
50    value: f64,
51    #[allow(dead_code)]
52    currency: String,
53}
54
55/// OpenAI usage API response structure.
56#[derive(serde::Deserialize, Debug)]
57struct OpenAIUsageResponse {
58    data: Vec<OpenAIUsageItem>,
59}
60
61#[derive(serde::Deserialize, Debug)]
62struct OpenAIUsageItem {
63    #[allow(dead_code)]
64    model: Option<String>,
65    #[allow(dead_code)]
66    num_requests: Option<u64>,
67    input_tokens: Option<u64>,
68    output_tokens: Option<u64>,
69    #[allow(dead_code)]
70    cached_input_tokens: Option<u64>,
71}
72
73/// OpenAI provider.
74pub struct OpenAIProvider {
75    metadata: ProviderMetadata,
76    /// Base URL override for tests.
77    base_url: Option<String>,
78}
79
80impl OpenAIProvider {
81    pub fn new() -> Self {
82        Self {
83            metadata: ProviderMetadata {
84                id: "openai",
85                name: "OpenAI",
86                description: "OpenAI API usage monitor",
87                auth_methods: &["api_key", "env"],
88                website: Some("https://platform.openai.com"),
89            },
90            base_url: None,
91        }
92    }
93
94    /// Creates a provider with a custom base URL (for tests).
95    pub fn with_base_url(url: &str) -> Self {
96        let mut p = Self::new();
97        p.base_url = Some(url.to_string());
98        p
99    }
100
101    fn api_base(&self) -> &str {
102        self.base_url.as_deref().unwrap_or("https://api.openai.com")
103    }
104
105    /// Detection helper: a non-empty API key is available.
106    fn detect_credentials_from(key: Option<&str>) -> bool {
107        key.is_some_and(|k| !k.is_empty())
108    }
109
110    /// Extracts the API key from context or environment variable.
111    fn resolve_api_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
112        if let Some(key) = ctx.config.get("api_key")
113            && !key.is_empty()
114        {
115            return Ok(key.clone());
116        }
117        if let Ok(key) = std::env::var("OPENAI_API_KEY")
118            && !key.is_empty()
119        {
120            return Ok(key);
121        }
122        Err(SpendPanelError::AuthFailed(
123            "openai".into(),
124            "no API key found in config or OPENAI_API_KEY env var".into(),
125        ))
126    }
127
128    /// Fetches rate limits (available in response headers).
129    async fn fetch_rate_limits(
130        base_url: &str,
131        client: &reqwest::Client,
132        api_key: &str,
133    ) -> Result<OpenAIRateLimitHeaders, SpendPanelError> {
134        // Light request just to capture rate limit headers
135        let resp = client
136            .get(format!("{}/v1/models", base_url))
137            .header("Authorization", format!("Bearer {}", api_key))
138            .send()
139            .await
140            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
141
142        if resp.status().is_success() || resp.status().as_u16() == 429 {
143            Ok(OpenAIRateLimitHeaders::from_headers(resp.headers()))
144        } else if resp.status().as_u16() == 401 {
145            Err(SpendPanelError::AuthFailed(
146                "openai".into(),
147                "invalid API key".into(),
148            ))
149        } else {
150            Err(SpendPanelError::ProviderError(
151                "openai".into(),
152                format!("unexpected status: {}", resp.status()),
153            ))
154        }
155    }
156
157    /// Fetches organization costs (requires admin key; failures become an empty list).
158    async fn fetch_costs(
159        base_url: &str,
160        client: &reqwest::Client,
161        api_key: &str,
162    ) -> Result<Vec<DailyCost>, SpendPanelError> {
163        let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
164        let week_ago = (chrono::Utc::now() - chrono::Duration::days(7))
165            .format("%Y-%m-%d")
166            .to_string();
167
168        let resp = client
169            .get(format!("{}/v1/organization/costs", base_url))
170            .query(&[("start_time", &*week_ago), ("end_time", &*today)])
171            .header("Authorization", format!("Bearer {}", api_key))
172            .send()
173            .await
174            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
175
176        if !resp.status().is_success() {
177            // Admin key required for costs; fail silently
178            return Ok(Vec::new());
179        }
180
181        let body = resp
182            .text()
183            .await
184            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
185
186        // Try to parse, but fail silently if not admin
187        if let Ok(cost_resp) = serde_json::from_str::<OpenAICostResponse>(&body) {
188            // Group by date (simplified: API does not return per-item dates yet)
189            let mut daily = std::collections::HashMap::new();
190            for item in cost_resp.data {
191                let date = chrono::Utc::now().date_naive();
192                let entry = daily.entry(date).or_insert(DailyCost {
193                    date,
194                    cost: 0.0,
195                    tokens_input: None,
196                    tokens_output: None,
197                    requests: None,
198                });
199                entry.cost += item.amount.value;
200            }
201            Ok(daily.into_values().collect())
202        } else {
203            Ok(Vec::new())
204        }
205    }
206
207    /// Fetches token usage (requires admin key; failures become an empty list).
208    async fn fetch_usage_report(
209        base_url: &str,
210        client: &reqwest::Client,
211        api_key: &str,
212    ) -> Result<Vec<DailyCost>, SpendPanelError> {
213        let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
214        let week_ago = (chrono::Utc::now() - chrono::Duration::days(7))
215            .format("%Y-%m-%d")
216            .to_string();
217
218        let resp = client
219            .get(format!("{}/v1/organization/usage/completions", base_url))
220            .query(&[
221                ("start_time", &*week_ago),
222                ("end_time", &*today),
223                ("bucket_width", "1d"),
224            ])
225            .header("Authorization", format!("Bearer {}", api_key))
226            .send()
227            .await
228            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
229
230        if !resp.status().is_success() {
231            return Ok(Vec::new());
232        }
233
234        let body = resp
235            .text()
236            .await
237            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
238
239        if let Ok(usage_resp) = serde_json::from_str::<OpenAIUsageResponse>(&body) {
240            // Group by date
241            let mut daily = std::collections::HashMap::new();
242            let today_naive = chrono::Utc::now().date_naive();
243            for item in usage_resp.data {
244                let date = today_naive;
245                let entry = daily.entry(date).or_insert(DailyCost {
246                    date,
247                    cost: 0.0,
248                    tokens_input: None,
249                    tokens_output: None,
250                    requests: None,
251                });
252                entry.tokens_input =
253                    Some(entry.tokens_input.unwrap_or(0) + item.input_tokens.unwrap_or(0));
254                entry.tokens_output =
255                    Some(entry.tokens_output.unwrap_or(0) + item.output_tokens.unwrap_or(0));
256            }
257            Ok(daily.into_values().collect())
258        } else {
259            Ok(Vec::new())
260        }
261    }
262}
263
264impl Default for OpenAIProvider {
265    fn default() -> Self {
266        Self::new()
267    }
268}
269
270#[async_trait]
271impl UsageProvider for OpenAIProvider {
272    fn metadata(&self) -> &ProviderMetadata {
273        &self.metadata
274    }
275
276    fn detect_credentials(&self) -> bool {
277        OpenAIProvider::detect_credentials_from(std::env::var("OPENAI_API_KEY").ok().as_deref())
278    }
279
280    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
281        let api_key = Self::resolve_api_key(ctx)?;
282        let base = self.api_base();
283
284        let client = reqwest::Client::builder()
285            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
286            .build()
287            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
288
289        // Fetch rate limits, costs, and usage in parallel
290        let (rate_limits, costs, usage) = tokio::join!(
291            Self::fetch_rate_limits(base, &client, &api_key),
292            Self::fetch_costs(base, &client, &api_key),
293            Self::fetch_usage_report(base, &client, &api_key),
294        );
295
296        let mut snapshot = UsageSnapshot::new("openai");
297        snapshot.collected_at = Utc::now();
298
299        // Rate limits
300        if let Ok(rl) = rate_limits {
301            if let (Some(limit), Some(remaining)) = (rl.limit_requests, rl.remaining_requests) {
302                let used = limit.saturating_sub(remaining);
303                snapshot.primary_rate_window = Some(RateWindow::new(used, limit, "RPM", 1));
304            }
305            if let (Some(limit), Some(remaining)) = (rl.limit_tokens, rl.remaining_tokens) {
306                let used = limit.saturating_sub(remaining);
307                snapshot.secondary_rate_window = Some(RateWindow::new(used, limit, "TPM", 1));
308            }
309        }
310
311        // Costs (with usage report tokens merged by date)
312        if let Ok(mut cost_list) = costs {
313            if let Ok(usage_list) = usage {
314                for u in usage_list {
315                    if let Some(entry) = cost_list.iter_mut().find(|c| c.date == u.date) {
316                        entry.tokens_input = u.tokens_input;
317                        entry.tokens_output = u.tokens_output;
318                    } else {
319                        cost_list.push(u);
320                    }
321                }
322            }
323
324            if !cost_list.is_empty() {
325                let total_cost: f64 = cost_list.iter().map(|c| c.cost).sum();
326
327                snapshot.cost = Some(CostSnapshot {
328                    total_cost: Some(total_cost),
329                    currency: "USD".into(),
330                    daily_costs: cost_list,
331                    spend_limit: Some(SpendLimit {
332                        limit: 100.0, // Tier 2 default
333                        used: total_cost,
334                        period: "monthly".into(),
335                    }),
336                });
337
338                snapshot.plan = Some(PlanInfo {
339                    name: "API".into(),
340                    tier: None,
341                    features: vec![],
342                    price: None,
343                    currency: Some("USD".into()),
344                    billing_period: Some("monthly".into()),
345                });
346            }
347        }
348
349        Ok(snapshot)
350    }
351}
352
353#[cfg(test)]
354mod tests {
355    use super::*;
356    use wiremock::matchers::{header, method, path};
357    use wiremock::{Mock, MockServer, ResponseTemplate};
358
359    /// Builds mocked rate limit headers.
360    fn rate_limit_headers() -> reqwest::header::HeaderMap {
361        let mut headers = reqwest::header::HeaderMap::new();
362        headers.insert("x-ratelimit-limit-requests", "100".parse().unwrap());
363        headers.insert("x-ratelimit-remaining-requests", "55".parse().unwrap());
364        headers.insert("x-ratelimit-limit-tokens", "40000".parse().unwrap());
365        headers.insert("x-ratelimit-remaining-tokens", "38800".parse().unwrap());
366        headers
367    }
368
369    #[test]
370    fn test_resolve_api_key_from_context() {
371        let mut ctx = ProviderContext::new();
372        ctx.config.insert("api_key".into(), "sk-test-123".into());
373        assert_eq!(
374            OpenAIProvider::resolve_api_key(&ctx).unwrap(),
375            "sk-test-123"
376        );
377    }
378
379    #[test]
380    fn test_resolve_api_key_empty_context_fails() {
381        let ctx = ProviderContext::new();
382        let result = OpenAIProvider::resolve_api_key(&ctx);
383        assert!(matches!(result, Err(SpendPanelError::AuthFailed(_, _))));
384    }
385
386    #[test]
387    fn test_detect_credentials_from() {
388        assert!(OpenAIProvider::detect_credentials_from(Some("sk-x")));
389        assert!(!OpenAIProvider::detect_credentials_from(Some("")));
390        assert!(!OpenAIProvider::detect_credentials_from(None));
391    }
392
393    #[test]
394    fn test_rate_limit_headers_parsing() {
395        let rl = OpenAIRateLimitHeaders::from_headers(&rate_limit_headers());
396        assert_eq!(rl.limit_requests, Some(100));
397        assert_eq!(rl.remaining_requests, Some(55));
398        assert_eq!(rl.limit_tokens, Some(40000));
399        assert_eq!(rl.remaining_tokens, Some(38800));
400    }
401
402    #[test]
403    fn test_rate_limit_headers_empty() {
404        let headers = reqwest::header::HeaderMap::new();
405        let rl = OpenAIRateLimitHeaders::from_headers(&headers);
406        assert!(rl.limit_requests.is_none());
407    }
408
409    #[tokio::test]
410    async fn test_fetch_rate_limits_success() {
411        let server = MockServer::start().await;
412
413        Mock::given(method("GET"))
414            .and(path("/v1/models"))
415            .and(header("Authorization", "Bearer sk-test"))
416            .respond_with(
417                ResponseTemplate::new(200)
418                    .set_body_json(serde_json::json!({"data": []}))
419                    .insert_header("x-ratelimit-limit-requests", "100")
420                    .insert_header("x-ratelimit-remaining-requests", "55"),
421            )
422            .mount(&server)
423            .await;
424
425        let client = reqwest::Client::new();
426        let rl = OpenAIProvider::fetch_rate_limits(&server.uri(), &client, "sk-test")
427            .await
428            .unwrap();
429        assert_eq!(rl.limit_requests, Some(100));
430        assert_eq!(rl.remaining_requests, Some(55));
431    }
432
433    #[tokio::test]
434    async fn test_fetch_rate_limits_401() {
435        let server = MockServer::start().await;
436        Mock::given(method("GET"))
437            .and(path("/v1/models"))
438            .respond_with(ResponseTemplate::new(401))
439            .mount(&server)
440            .await;
441
442        let client = reqwest::Client::new();
443        let result = OpenAIProvider::fetch_rate_limits(&server.uri(), &client, "bad").await;
444        assert!(matches!(result, Err(SpendPanelError::AuthFailed(_, _))));
445    }
446
447    #[tokio::test]
448    async fn test_fetch_costs_success() {
449        let server = MockServer::start().await;
450
451        Mock::given(method("GET"))
452            .and(path("/v1/organization/costs"))
453            .and(header("Authorization", "Bearer sk-test"))
454            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
455                "data": [{
456                    "amount": { "value": 1.20, "currency": "usd" },
457                    "line_item": "gpt-4o"
458                }]
459            })))
460            .mount(&server)
461            .await;
462
463        let client = reqwest::Client::new();
464        let costs = OpenAIProvider::fetch_costs(&server.uri(), &client, "sk-test")
465            .await
466            .unwrap();
467        assert_eq!(costs.len(), 1);
468        assert!((costs[0].cost - 1.20).abs() < f64::EPSILON);
469    }
470
471    #[tokio::test]
472    async fn test_fetch_costs_non_admin_returns_empty() {
473        let server = MockServer::start().await;
474        Mock::given(method("GET"))
475            .and(path("/v1/organization/costs"))
476            .respond_with(ResponseTemplate::new(401))
477            .mount(&server)
478            .await;
479
480        let client = reqwest::Client::new();
481        let costs = OpenAIProvider::fetch_costs(&server.uri(), &client, "non-admin")
482            .await
483            .unwrap();
484        assert!(costs.is_empty());
485    }
486
487    #[tokio::test]
488    async fn test_integrated_fetch_with_mocks() {
489        let server = MockServer::start().await;
490
491        Mock::given(method("GET"))
492            .and(path("/v1/models"))
493            .respond_with(
494                ResponseTemplate::new(200)
495                    .set_body_json(serde_json::json!({"data": []}))
496                    .insert_header("x-ratelimit-limit-requests", "100")
497                    .insert_header("x-ratelimit-remaining-requests", "55")
498                    .insert_header("x-ratelimit-limit-tokens", "40000")
499                    .insert_header("x-ratelimit-remaining-tokens", "38800"),
500            )
501            .mount(&server)
502            .await;
503
504        Mock::given(method("GET"))
505            .and(path("/v1/organization/costs"))
506            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
507                "data": [{"amount": {"value": 2.40, "currency": "usd"}, "line_item": "gpt-4o"}]
508            })))
509            .mount(&server)
510            .await;
511
512        Mock::given(method("GET"))
513            .and(path("/v1/organization/usage/completions"))
514            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
515                "data": [{"model": "gpt-4o", "num_requests": 450, "input_tokens": 45000, "output_tokens": 12000}]
516            })))
517            .mount(&server)
518            .await;
519
520        let provider = OpenAIProvider::with_base_url(&server.uri());
521        let mut ctx = ProviderContext::new();
522        ctx.config.insert("api_key".into(), "sk-test".into());
523
524        let snap = provider.fetch_usage(&ctx).await.unwrap();
525        assert_eq!(snap.provider_id, "openai");
526
527        let primary = snap.primary_rate_window.unwrap();
528        assert_eq!(primary.used, Some(45)); // 100 - 55
529
530        let secondary = snap.secondary_rate_window.unwrap();
531        assert_eq!(secondary.used, Some(1200)); // 40000 - 38800
532
533        let cost = snap.cost.unwrap();
534        assert!((cost.total_cost.unwrap() - 2.40).abs() < f64::EPSILON);
535        assert_eq!(cost.daily_costs[0].tokens_input, Some(45000));
536    }
537}
usage_monitor_cli/provider/openai.rs

usage_monitor_cli/provider/
openai.rs
