openai_ergonomic/builders/
batch.rs

1//! Batch API builders.
2//!
3//! This module provides ergonomic builders for `OpenAI` Batch API operations,
4//! which allow you to send asynchronous groups of requests with 24-hour turnaround
5//! and 50% cost reduction compared to synchronous API calls.
6
7use std::collections::HashMap;
8
9/// Builder for creating batch jobs.
10///
11/// Batch jobs allow you to process multiple requests asynchronously at a lower cost
12/// with longer processing time (up to 24 hours).
13#[derive(Debug, Clone)]
14pub struct BatchJobBuilder {
15    input_file_id: String,
16    endpoint: BatchEndpoint,
17    completion_window: BatchCompletionWindow,
18    metadata: HashMap<String, String>,
19}
20
21/// Supported endpoints for batch processing.
22#[derive(Debug, Clone)]
23pub enum BatchEndpoint {
24    /// Chat completions endpoint
25    ChatCompletions,
26    /// Embeddings endpoint
27    Embeddings,
28    /// Completions endpoint (legacy)
29    Completions,
30}
31
32/// Completion window for batch jobs.
33#[derive(Debug, Clone)]
34pub enum BatchCompletionWindow {
35    /// Complete within 24 hours
36    Hours24,
37}
38
39impl BatchJobBuilder {
40    /// Create a new batch job builder.
41    ///
42    /// # Examples
43    ///
44    /// ```rust
45    /// use openai_ergonomic::builders::batch::{BatchJobBuilder, BatchEndpoint};
46    ///
47    /// let builder = BatchJobBuilder::new("file-batch-input", BatchEndpoint::ChatCompletions);
48    /// ```
49    #[must_use]
50    pub fn new(input_file_id: impl Into<String>, endpoint: BatchEndpoint) -> Self {
51        Self {
52            input_file_id: input_file_id.into(),
53            endpoint,
54            completion_window: BatchCompletionWindow::Hours24,
55            metadata: HashMap::new(),
56        }
57    }
58
59    /// Set the completion window for the batch job.
60    #[must_use]
61    pub fn completion_window(mut self, window: BatchCompletionWindow) -> Self {
62        self.completion_window = window;
63        self
64    }
65
66    /// Add metadata to the batch job.
67    #[must_use]
68    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
69        self.metadata.insert(key.into(), value.into());
70        self
71    }
72
73    /// Get the input file ID.
74    #[must_use]
75    pub fn input_file_id(&self) -> &str {
76        &self.input_file_id
77    }
78
79    /// Get the endpoint for this batch job.
80    #[must_use]
81    pub fn endpoint(&self) -> &BatchEndpoint {
82        &self.endpoint
83    }
84
85    /// Get the completion window.
86    #[must_use]
87    pub fn completion_window_ref(&self) -> &BatchCompletionWindow {
88        &self.completion_window
89    }
90
91    /// Get the metadata.
92    #[must_use]
93    pub fn metadata_ref(&self) -> &HashMap<String, String> {
94        &self.metadata
95    }
96
97    /// Check if metadata is empty.
98    #[must_use]
99    pub fn has_metadata(&self) -> bool {
100        !self.metadata.is_empty()
101    }
102}
103
104/// Builder for listing batch jobs.
105#[derive(Debug, Clone, Default)]
106pub struct BatchJobListBuilder {
107    after: Option<String>,
108    limit: Option<i32>,
109}
110
111impl BatchJobListBuilder {
112    /// Create a new batch job list builder.
113    #[must_use]
114    pub fn new() -> Self {
115        Self::default()
116    }
117
118    /// Set the cursor for pagination.
119    #[must_use]
120    pub fn after(mut self, cursor: impl Into<String>) -> Self {
121        self.after = Some(cursor.into());
122        self
123    }
124
125    /// Set the maximum number of jobs to return.
126    #[must_use]
127    pub fn limit(mut self, limit: i32) -> Self {
128        self.limit = Some(limit);
129        self
130    }
131
132    /// Get the pagination cursor.
133    #[must_use]
134    pub fn after_ref(&self) -> Option<&str> {
135        self.after.as_deref()
136    }
137
138    /// Get the limit.
139    #[must_use]
140    pub fn limit_ref(&self) -> Option<i32> {
141        self.limit
142    }
143}
144
145/// Builder for retrieving batch job details.
146#[derive(Debug, Clone)]
147pub struct BatchJobRetrievalBuilder {
148    batch_id: String,
149}
150
151impl BatchJobRetrievalBuilder {
152    /// Create a new batch job retrieval builder.
153    #[must_use]
154    pub fn new(batch_id: impl Into<String>) -> Self {
155        Self {
156            batch_id: batch_id.into(),
157        }
158    }
159
160    /// Get the batch ID.
161    #[must_use]
162    pub fn batch_id(&self) -> &str {
163        &self.batch_id
164    }
165}
166
167/// Builder for cancelling batch jobs.
168#[derive(Debug, Clone)]
169pub struct BatchJobCancelBuilder {
170    batch_id: String,
171}
172
173impl BatchJobCancelBuilder {
174    /// Create a new batch job cancel builder.
175    #[must_use]
176    pub fn new(batch_id: impl Into<String>) -> Self {
177        Self {
178            batch_id: batch_id.into(),
179        }
180    }
181
182    /// Get the batch ID.
183    #[must_use]
184    pub fn batch_id(&self) -> &str {
185        &self.batch_id
186    }
187}
188
189/// Helper function to create a chat completions batch job.
190#[must_use]
191pub fn batch_chat_completions(input_file_id: impl Into<String>) -> BatchJobBuilder {
192    BatchJobBuilder::new(input_file_id, BatchEndpoint::ChatCompletions)
193}
194
195/// Helper function to create an embeddings batch job.
196#[must_use]
197pub fn batch_embeddings(input_file_id: impl Into<String>) -> BatchJobBuilder {
198    BatchJobBuilder::new(input_file_id, BatchEndpoint::Embeddings)
199}
200
201/// Helper function to create a completions batch job.
202#[must_use]
203pub fn batch_completions(input_file_id: impl Into<String>) -> BatchJobBuilder {
204    BatchJobBuilder::new(input_file_id, BatchEndpoint::Completions)
205}
206
207/// Helper function to create a batch job with metadata.
208#[must_use]
209#[allow(clippy::implicit_hasher)]
210pub fn batch_job_with_metadata(
211    input_file_id: impl Into<String>,
212    endpoint: BatchEndpoint,
213    metadata: HashMap<String, String>,
214) -> BatchJobBuilder {
215    let mut builder = BatchJobBuilder::new(input_file_id, endpoint);
216    for (key, value) in metadata {
217        builder = builder.metadata(key, value);
218    }
219    builder
220}
221
222/// Helper function to list batch jobs.
223#[must_use]
224pub fn list_batch_jobs() -> BatchJobListBuilder {
225    BatchJobListBuilder::new()
226}
227
228/// Helper function to retrieve a specific batch job.
229#[must_use]
230pub fn get_batch_job(batch_id: impl Into<String>) -> BatchJobRetrievalBuilder {
231    BatchJobRetrievalBuilder::new(batch_id)
232}
233
234/// Helper function to cancel a batch job.
235#[must_use]
236pub fn cancel_batch_job(batch_id: impl Into<String>) -> BatchJobCancelBuilder {
237    BatchJobCancelBuilder::new(batch_id)
238}
239
240impl std::fmt::Display for BatchEndpoint {
241    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        match self {
243            BatchEndpoint::ChatCompletions => write!(f, "/v1/chat/completions"),
244            BatchEndpoint::Embeddings => write!(f, "/v1/embeddings"),
245            BatchEndpoint::Completions => write!(f, "/v1/completions"),
246        }
247    }
248}
249
250impl std::fmt::Display for BatchCompletionWindow {
251    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
252        match self {
253            BatchCompletionWindow::Hours24 => write!(f, "24h"),
254        }
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261
262    #[test]
263    fn test_batch_job_builder_new() {
264        let builder = BatchJobBuilder::new("file-input", BatchEndpoint::ChatCompletions);
265
266        assert_eq!(builder.input_file_id(), "file-input");
267        match builder.endpoint() {
268            BatchEndpoint::ChatCompletions => {}
269            _ => panic!("Expected ChatCompletions endpoint"),
270        }
271        assert!(!builder.has_metadata());
272    }
273
274    #[test]
275    fn test_batch_job_builder_with_metadata() {
276        let builder = BatchJobBuilder::new("file-input", BatchEndpoint::Embeddings)
277            .metadata("project", "test-project")
278            .metadata("version", "v1");
279
280        assert!(builder.has_metadata());
281        assert_eq!(builder.metadata_ref().len(), 2);
282        assert_eq!(
283            builder.metadata_ref().get("project"),
284            Some(&"test-project".to_string())
285        );
286        assert_eq!(
287            builder.metadata_ref().get("version"),
288            Some(&"v1".to_string())
289        );
290    }
291
292    #[test]
293    fn test_batch_job_builder_completion_window() {
294        let builder = BatchJobBuilder::new("file-input", BatchEndpoint::ChatCompletions)
295            .completion_window(BatchCompletionWindow::Hours24);
296
297        match builder.completion_window_ref() {
298            BatchCompletionWindow::Hours24 => {}
299        }
300    }
301
302    #[test]
303    fn test_batch_job_list_builder() {
304        let builder = BatchJobListBuilder::new().after("batch-123").limit(10);
305
306        assert_eq!(builder.after_ref(), Some("batch-123"));
307        assert_eq!(builder.limit_ref(), Some(10));
308    }
309
310    #[test]
311    fn test_batch_job_retrieval_builder() {
312        let builder = BatchJobRetrievalBuilder::new("batch-456");
313        assert_eq!(builder.batch_id(), "batch-456");
314    }
315
316    #[test]
317    fn test_batch_job_cancel_builder() {
318        let builder = BatchJobCancelBuilder::new("batch-789");
319        assert_eq!(builder.batch_id(), "batch-789");
320    }
321
322    #[test]
323    fn test_batch_chat_completions_helper() {
324        let builder = batch_chat_completions("file-input");
325        assert_eq!(builder.input_file_id(), "file-input");
326        match builder.endpoint() {
327            BatchEndpoint::ChatCompletions => {}
328            _ => panic!("Expected ChatCompletions endpoint"),
329        }
330    }
331
332    #[test]
333    fn test_batch_embeddings_helper() {
334        let builder = batch_embeddings("file-input");
335        match builder.endpoint() {
336            BatchEndpoint::Embeddings => {}
337            _ => panic!("Expected Embeddings endpoint"),
338        }
339    }
340
341    #[test]
342    fn test_batch_completions_helper() {
343        let builder = batch_completions("file-input");
344        match builder.endpoint() {
345            BatchEndpoint::Completions => {}
346            _ => panic!("Expected Completions endpoint"),
347        }
348    }
349
350    #[test]
351    fn test_batch_job_with_metadata_helper() {
352        let mut metadata = HashMap::new();
353        metadata.insert("key1".to_string(), "value1".to_string());
354        metadata.insert("key2".to_string(), "value2".to_string());
355
356        let builder =
357            batch_job_with_metadata("file-input", BatchEndpoint::ChatCompletions, metadata);
358
359        assert!(builder.has_metadata());
360        assert_eq!(builder.metadata_ref().len(), 2);
361    }
362
363    #[test]
364    fn test_list_batch_jobs_helper() {
365        let builder = list_batch_jobs();
366        assert!(builder.after_ref().is_none());
367        assert!(builder.limit_ref().is_none());
368    }
369
370    #[test]
371    fn test_get_batch_job_helper() {
372        let builder = get_batch_job("batch-123");
373        assert_eq!(builder.batch_id(), "batch-123");
374    }
375
376    #[test]
377    fn test_cancel_batch_job_helper() {
378        let builder = cancel_batch_job("batch-456");
379        assert_eq!(builder.batch_id(), "batch-456");
380    }
381
382    #[test]
383    fn test_batch_endpoint_display() {
384        assert_eq!(
385            BatchEndpoint::ChatCompletions.to_string(),
386            "/v1/chat/completions"
387        );
388        assert_eq!(BatchEndpoint::Embeddings.to_string(), "/v1/embeddings");
389        assert_eq!(BatchEndpoint::Completions.to_string(), "/v1/completions");
390    }
391
392    #[test]
393    fn test_batch_completion_window_display() {
394        assert_eq!(BatchCompletionWindow::Hours24.to_string(), "24h");
395    }
396
397    #[test]
398    fn test_batch_job_list_builder_default() {
399        let builder = BatchJobListBuilder::default();
400        assert!(builder.after_ref().is_none());
401        assert!(builder.limit_ref().is_none());
402    }
403}
openai_ergonomic/builders/batch.rs

openai_ergonomic/builders/
batch.rs
