gemini_rust/batch/
builder.rs

1use snafu::ResultExt;
2use std::sync::Arc;
3
4use super::handle::BatchHandle;
5use super::model::*;
6use super::*;
7use crate::{client::GeminiClient, generation::GenerateContentRequest};
8
9/// A builder for creating and executing synchronous batch content generation requests.
10///
11/// This builder simplifies the process of constructing a batch request, allowing you to
12/// add multiple `GenerateContentRequest` items and then execute them as a single
13/// long-running operation.
14pub struct BatchBuilder {
15    client: Arc<GeminiClient>,
16    display_name: String,
17    requests: Vec<GenerateContentRequest>,
18}
19
20impl BatchBuilder {
21    /// Create a new batch builder
22    pub(crate) fn new(client: Arc<GeminiClient>) -> Self {
23        Self {
24            client,
25            display_name: "RustBatch".to_string(),
26            requests: Vec::new(),
27        }
28    }
29
30    /// Sets the user-friendly display name for the batch request.
31    pub fn with_name(mut self, name: String) -> Self {
32        self.display_name = name;
33        self
34    }
35
36    /// Sets all requests for the batch operation, replacing any existing requests.
37    pub fn with_requests(mut self, requests: Vec<GenerateContentRequest>) -> Self {
38        self.requests = requests;
39        self
40    }
41
42    /// Adds a single `GenerateContentRequest` to the batch.
43    pub fn with_request(mut self, request: GenerateContentRequest) -> Self {
44        self.requests.push(request);
45        self
46    }
47
48    /// Constructs the final `BatchGenerateContentRequest` from the builder's configuration.
49    ///
50    /// This method consumes the builder.
51    pub fn build(self) -> BatchGenerateContentRequest {
52        let batch_requests: Vec<BatchRequestItem> = self
53            .requests
54            .into_iter()
55            .enumerate()
56            .map(|(key, request)| BatchRequestItem {
57                request,
58                metadata: RequestMetadata { key },
59            })
60            .collect();
61
62        BatchGenerateContentRequest {
63            batch: BatchConfig {
64                display_name: self.display_name,
65                input_config: InputConfig::Requests(RequestsContainer {
66                    requests: batch_requests,
67                }),
68            },
69        }
70    }
71
72    /// Submits the batch request to the Gemini API and returns a `Batch` handle.
73    ///
74    /// This method consumes the builder and initiates the long-running batch operation.
75    pub async fn execute(self) -> Result<BatchHandle, Error> {
76        let client = self.client.clone();
77        let request = self.build();
78        let response = client
79            .batch_generate_content_sync(request)
80            .await
81            .context(ClientSnafu)?;
82        Ok(BatchHandle::new(response.name, client))
83    }
84
85    /// Executes the batch request by first uploading the requests as a JSON file.
86    ///
87    /// This method is ideal for large batch jobs that might exceed inline request limits.
88    /// It consumes the builder, serializes the requests to the JSON Lines format,
89    /// uploads the content as a file, and then starts the batch operation using that file.
90    pub async fn execute_as_file(self) -> Result<BatchHandle, Error> {
91        let mut json_lines = String::new();
92        for (index, item) in self.requests.into_iter().enumerate() {
93            let item = BatchRequestFileItem {
94                request: item,
95                key: index,
96            };
97
98            let line = serde_json::to_string(&item).context(SerializeSnafu)?;
99            json_lines.push_str(&line);
100            json_lines.push('\n');
101        }
102        let json_bytes = json_lines.into_bytes();
103
104        let file_display_name = format!("{}-input.jsonl", self.display_name);
105        let file = crate::files::builder::FileBuilder::new(self.client.clone(), json_bytes)
106            .display_name(file_display_name)
107            .with_mime_type(
108                "application/jsonl"
109                    .parse()
110                    .expect("failed to parse MIME type 'application/jsonl'"),
111            )
112            .upload()
113            .await
114            .context(FileSnafu)?;
115
116        let request = BatchGenerateContentRequest {
117            batch: BatchConfig {
118                display_name: self.display_name,
119                input_config: InputConfig::FileName(file.name().to_string()),
120            },
121        };
122
123        let client = self.client.clone();
124        let response = client
125            .batch_generate_content_sync(request)
126            .await
127            .context(ClientSnafu)?;
128
129        Ok(BatchHandle::new(response.name, client))
130    }
131}
gemini_rust/batch/builder.rs

gemini_rust/batch/
builder.rs
