gemini_rust/
batch_builder.rs

1use std::sync::Arc;
2
3use crate::{
4    batch::Batch,
5    client::GeminiClient,
6    models::{
7        BatchConfig, BatchGenerateContentRequest, BatchRequestItem, GenerateContentRequest,
8        InputConfig, RequestMetadata, RequestsContainer,
9    },
10    Result,
11};
12
13/// A builder for creating and executing synchronous batch content generation requests.
14///
15/// This builder simplifies the process of constructing a batch request, allowing you to
16/// add multiple `GenerateContentRequest` items and then execute them as a single
17/// long-running operation.
18pub struct BatchBuilder {
19    client: Arc<GeminiClient>,
20    display_name: String,
21    requests: Vec<GenerateContentRequest>,
22}
23
24impl BatchBuilder {
25    /// Create a new batch builder
26    pub(crate) fn new(client: Arc<GeminiClient>) -> Self {
27        Self {
28            client,
29            display_name: "RustBatch".to_string(),
30            requests: Vec::new(),
31        }
32    }
33
34    /// Sets the user-friendly display name for the batch request.
35    pub fn with_name(mut self, name: String) -> Self {
36        self.display_name = name;
37        self
38    }
39
40    /// Sets all requests for the batch operation, replacing any existing requests.
41    pub fn with_requests(mut self, requests: Vec<GenerateContentRequest>) -> Self {
42        self.requests = requests;
43        self
44    }
45
46    /// Adds a single `GenerateContentRequest` to the batch.
47    pub fn with_request(mut self, request: GenerateContentRequest) -> Self {
48        self.requests.push(request);
49        self
50    }
51
52    /// Constructs the final `BatchGenerateContentRequest` from the builder's configuration.
53    ///
54    /// This method consumes the builder.
55    pub fn build(self) -> BatchGenerateContentRequest {
56        let batch_requests: Vec<BatchRequestItem> = self
57            .requests
58            .into_iter()
59            .enumerate()
60            .map(|(i, request)| BatchRequestItem {
61                request,
62                metadata: Some(RequestMetadata { key: i.to_string() }),
63            })
64            .collect();
65
66        BatchGenerateContentRequest {
67            batch: BatchConfig {
68                display_name: self.display_name,
69                input_config: InputConfig {
70                    requests: RequestsContainer {
71                        requests: batch_requests,
72                    },
73                },
74            },
75        }
76    }
77
78    /// Submits the batch request to the Gemini API and returns a `Batch` handle.
79    ///
80    /// This method consumes the builder and initiates the long-running batch operation.
81    pub async fn execute(self) -> Result<Batch> {
82        let client = self.client.clone();
83        let request = self.build();
84        let response = client.batch_generate_content_sync(request).await?;
85        Ok(Batch::new(response.name, client))
86    }
87}
gemini_rust/batch_builder.rs

gemini_rust/
batch_builder.rs
