openai_ergonomic/builders/
usage.rs

1//! Usage API builders.
2//!
3//! Provides high-level builders for querying usage and cost data from the `OpenAI` API.
4//! Supports filtering by date range, aggregation buckets, projects, users, API keys, and models.
5
6/// Bucket width for aggregating usage data.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum BucketWidth {
9    /// Aggregate data by day.
10    Day,
11    /// Aggregate data by hour.
12    Hour,
13}
14
15impl BucketWidth {
16    /// Convert to API string representation.
17    #[must_use]
18    pub fn as_str(&self) -> &'static str {
19        match self {
20            Self::Day => "1d",
21            Self::Hour => "1h",
22        }
23    }
24}
25
26impl std::fmt::Display for BucketWidth {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        write!(f, "{}", self.as_str())
29    }
30}
31
32/// Group by field for usage aggregation.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum GroupBy {
35    /// Group by project ID.
36    ProjectId,
37    /// Group by user ID.
38    UserId,
39    /// Group by API key ID.
40    ApiKeyId,
41    /// Group by model.
42    Model,
43}
44
45impl GroupBy {
46    /// Convert to API string representation.
47    #[must_use]
48    pub fn as_str(&self) -> &'static str {
49        match self {
50            Self::ProjectId => "project_id",
51            Self::UserId => "user_id",
52            Self::ApiKeyId => "api_key_id",
53            Self::Model => "model",
54        }
55    }
56}
57
58impl std::fmt::Display for GroupBy {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        write!(f, "{}", self.as_str())
61    }
62}
63
64/// Builder for querying usage data from the `OpenAI` API.
65///
66/// # Examples
67///
68/// ```rust
69/// use openai_ergonomic::builders::usage::{UsageBuilder, BucketWidth};
70///
71/// let builder = UsageBuilder::new(1704067200, None) // Start time (Unix timestamp)
72///     .bucket_width(BucketWidth::Day)
73///     .limit(100);
74/// ```
75#[derive(Debug, Clone)]
76pub struct UsageBuilder {
77    start_time: i32,
78    end_time: Option<i32>,
79    bucket_width: Option<BucketWidth>,
80    project_ids: Vec<String>,
81    user_ids: Vec<String>,
82    api_key_ids: Vec<String>,
83    models: Vec<String>,
84    group_by: Vec<GroupBy>,
85    limit: Option<i32>,
86    page: Option<String>,
87}
88
89impl UsageBuilder {
90    /// Create a new usage builder with the specified start time.
91    ///
92    /// # Arguments
93    ///
94    /// * `start_time` - Unix timestamp (in seconds) for the start of the query range
95    /// * `end_time` - Optional Unix timestamp (in seconds) for the end of the query range
96    #[must_use]
97    pub fn new(start_time: i32, end_time: Option<i32>) -> Self {
98        Self {
99            start_time,
100            end_time,
101            bucket_width: None,
102            project_ids: Vec::new(),
103            user_ids: Vec::new(),
104            api_key_ids: Vec::new(),
105            models: Vec::new(),
106            group_by: Vec::new(),
107            limit: None,
108            page: None,
109        }
110    }
111
112    /// Set the bucket width for aggregation.
113    #[must_use]
114    pub fn bucket_width(mut self, width: BucketWidth) -> Self {
115        self.bucket_width = Some(width);
116        self
117    }
118
119    /// Filter by a single project ID.
120    #[must_use]
121    pub fn project_id(mut self, id: impl Into<String>) -> Self {
122        self.project_ids.push(id.into());
123        self
124    }
125
126    /// Filter by multiple project IDs.
127    #[must_use]
128    pub fn project_ids<I, S>(mut self, ids: I) -> Self
129    where
130        I: IntoIterator<Item = S>,
131        S: Into<String>,
132    {
133        self.project_ids.extend(ids.into_iter().map(Into::into));
134        self
135    }
136
137    /// Filter by a single user ID.
138    #[must_use]
139    pub fn user_id(mut self, id: impl Into<String>) -> Self {
140        self.user_ids.push(id.into());
141        self
142    }
143
144    /// Filter by multiple user IDs.
145    #[must_use]
146    pub fn user_ids<I, S>(mut self, ids: I) -> Self
147    where
148        I: IntoIterator<Item = S>,
149        S: Into<String>,
150    {
151        self.user_ids.extend(ids.into_iter().map(Into::into));
152        self
153    }
154
155    /// Filter by a single API key ID.
156    #[must_use]
157    pub fn api_key_id(mut self, id: impl Into<String>) -> Self {
158        self.api_key_ids.push(id.into());
159        self
160    }
161
162    /// Filter by multiple API key IDs.
163    #[must_use]
164    pub fn api_key_ids<I, S>(mut self, ids: I) -> Self
165    where
166        I: IntoIterator<Item = S>,
167        S: Into<String>,
168    {
169        self.api_key_ids.extend(ids.into_iter().map(Into::into));
170        self
171    }
172
173    /// Filter by a single model.
174    #[must_use]
175    pub fn model(mut self, model: impl Into<String>) -> Self {
176        self.models.push(model.into());
177        self
178    }
179
180    /// Filter by multiple models.
181    #[must_use]
182    pub fn models<I, S>(mut self, models: I) -> Self
183    where
184        I: IntoIterator<Item = S>,
185        S: Into<String>,
186    {
187        self.models.extend(models.into_iter().map(Into::into));
188        self
189    }
190
191    /// Add a group by field.
192    #[must_use]
193    pub fn group_by(mut self, field: GroupBy) -> Self {
194        self.group_by.push(field);
195        self
196    }
197
198    /// Add multiple group by fields.
199    #[must_use]
200    pub fn group_by_fields<I>(mut self, fields: I) -> Self
201    where
202        I: IntoIterator<Item = GroupBy>,
203    {
204        self.group_by.extend(fields);
205        self
206    }
207
208    /// Set the maximum number of results to return.
209    #[must_use]
210    pub fn limit(mut self, limit: i32) -> Self {
211        self.limit = Some(limit);
212        self
213    }
214
215    /// Set the pagination cursor.
216    #[must_use]
217    pub fn page(mut self, page: impl Into<String>) -> Self {
218        self.page = Some(page.into());
219        self
220    }
221
222    /// Get the start time.
223    #[must_use]
224    pub fn start_time(&self) -> i32 {
225        self.start_time
226    }
227
228    /// Get the end time.
229    #[must_use]
230    pub fn end_time(&self) -> Option<i32> {
231        self.end_time
232    }
233
234    /// Get the bucket width.
235    #[must_use]
236    pub fn bucket_width_ref(&self) -> Option<BucketWidth> {
237        self.bucket_width
238    }
239
240    /// Get the project IDs.
241    #[must_use]
242    pub fn project_ids_ref(&self) -> &[String] {
243        &self.project_ids
244    }
245
246    /// Get the user IDs.
247    #[must_use]
248    pub fn user_ids_ref(&self) -> &[String] {
249        &self.user_ids
250    }
251
252    /// Get the API key IDs.
253    #[must_use]
254    pub fn api_key_ids_ref(&self) -> &[String] {
255        &self.api_key_ids
256    }
257
258    /// Get the models.
259    #[must_use]
260    pub fn models_ref(&self) -> &[String] {
261        &self.models
262    }
263
264    /// Get the group by fields.
265    #[must_use]
266    pub fn group_by_ref(&self) -> &[GroupBy] {
267        &self.group_by
268    }
269
270    /// Get the limit.
271    #[must_use]
272    pub fn limit_ref(&self) -> Option<i32> {
273        self.limit
274    }
275
276    /// Get the page cursor.
277    #[must_use]
278    pub fn page_ref(&self) -> Option<&str> {
279        self.page.as_deref()
280    }
281
282    /// Convert project IDs to `Option<Vec<String>>`.
283    #[must_use]
284    pub fn project_ids_option(&self) -> Option<Vec<String>> {
285        if self.project_ids.is_empty() {
286            None
287        } else {
288            Some(self.project_ids.clone())
289        }
290    }
291
292    /// Convert user IDs to `Option<Vec<String>>`.
293    #[must_use]
294    pub fn user_ids_option(&self) -> Option<Vec<String>> {
295        if self.user_ids.is_empty() {
296            None
297        } else {
298            Some(self.user_ids.clone())
299        }
300    }
301
302    /// Convert API key IDs to `Option<Vec<String>>`.
303    #[must_use]
304    pub fn api_key_ids_option(&self) -> Option<Vec<String>> {
305        if self.api_key_ids.is_empty() {
306            None
307        } else {
308            Some(self.api_key_ids.clone())
309        }
310    }
311
312    /// Convert models to `Option<Vec<String>>`.
313    #[must_use]
314    pub fn models_option(&self) -> Option<Vec<String>> {
315        if self.models.is_empty() {
316            None
317        } else {
318            Some(self.models.clone())
319        }
320    }
321
322    /// Convert group by fields to `Option<Vec<String>>`.
323    #[must_use]
324    pub fn group_by_option(&self) -> Option<Vec<String>> {
325        if self.group_by.is_empty() {
326            None
327        } else {
328            Some(self.group_by.iter().map(ToString::to_string).collect())
329        }
330    }
331
332    /// Get bucket width as `Option<&str>`.
333    #[must_use]
334    pub fn bucket_width_str(&self) -> Option<&str> {
335        self.bucket_width.as_ref().map(BucketWidth::as_str)
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn test_usage_builder_basic() {
345        let builder = UsageBuilder::new(1_704_067_200, None);
346        assert_eq!(builder.start_time(), 1_704_067_200);
347        assert_eq!(builder.end_time(), None);
348    }
349
350    #[test]
351    fn test_usage_builder_with_end_time() {
352        let builder = UsageBuilder::new(1_704_067_200, Some(1_704_153_600));
353        assert_eq!(builder.start_time(), 1_704_067_200);
354        assert_eq!(builder.end_time(), Some(1_704_153_600));
355    }
356
357    #[test]
358    fn test_usage_builder_with_bucket_width() {
359        let builder = UsageBuilder::new(1_704_067_200, None).bucket_width(BucketWidth::Day);
360        assert_eq!(builder.bucket_width_ref(), Some(BucketWidth::Day));
361        assert_eq!(builder.bucket_width_str(), Some("1d"));
362    }
363
364    #[test]
365    fn test_usage_builder_with_filters() {
366        let builder = UsageBuilder::new(1_704_067_200, None)
367            .project_id("proj_123")
368            .user_id("user_456")
369            .model("gpt-4");
370
371        assert_eq!(builder.project_ids_ref(), &["proj_123"]);
372        assert_eq!(builder.user_ids_ref(), &["user_456"]);
373        assert_eq!(builder.models_ref(), &["gpt-4"]);
374    }
375
376    #[test]
377    fn test_usage_builder_with_multiple_filters() {
378        let builder = UsageBuilder::new(1_704_067_200, None)
379            .project_ids(vec!["proj_1", "proj_2"])
380            .user_ids(vec!["user_1", "user_2"])
381            .models(vec!["gpt-4", "gpt-3.5-turbo"]);
382
383        assert_eq!(builder.project_ids_ref().len(), 2);
384        assert_eq!(builder.user_ids_ref().len(), 2);
385        assert_eq!(builder.models_ref().len(), 2);
386    }
387
388    #[test]
389    fn test_usage_builder_with_group_by() {
390        let builder = UsageBuilder::new(1_704_067_200, None)
391            .group_by(GroupBy::ProjectId)
392            .group_by(GroupBy::Model);
393
394        assert_eq!(builder.group_by_ref().len(), 2);
395        let group_by_strings = builder.group_by_option().unwrap();
396        assert_eq!(group_by_strings, vec!["project_id", "model"]);
397    }
398
399    #[test]
400    fn test_usage_builder_with_pagination() {
401        let builder = UsageBuilder::new(1_704_067_200, None)
402            .limit(50)
403            .page("next_page_token");
404
405        assert_eq!(builder.limit_ref(), Some(50));
406        assert_eq!(builder.page_ref(), Some("next_page_token"));
407    }
408
409    #[test]
410    fn test_bucket_width_display() {
411        assert_eq!(BucketWidth::Day.to_string(), "1d");
412        assert_eq!(BucketWidth::Hour.to_string(), "1h");
413    }
414
415    #[test]
416    fn test_group_by_display() {
417        assert_eq!(GroupBy::ProjectId.to_string(), "project_id");
418        assert_eq!(GroupBy::UserId.to_string(), "user_id");
419        assert_eq!(GroupBy::ApiKeyId.to_string(), "api_key_id");
420        assert_eq!(GroupBy::Model.to_string(), "model");
421    }
422
423    #[test]
424    fn test_empty_vectors_to_none() {
425        let builder = UsageBuilder::new(1_704_067_200, None);
426        assert!(builder.project_ids_option().is_none());
427        assert!(builder.user_ids_option().is_none());
428        assert!(builder.api_key_ids_option().is_none());
429        assert!(builder.models_option().is_none());
430        assert!(builder.group_by_option().is_none());
431    }
432}
openai_ergonomic/builders/usage.rs

openai_ergonomic/builders/
usage.rs
