gemini_rust/
batch_builder.rs

1use std::sync::Arc;
2
3use crate::{
4    client::GeminiClient,
5    models::{
6        BatchConfig, BatchGenerateContentRequest, BatchGenerateContentResponse, BatchRequestItem,
7        GenerateContentRequest, InputConfig, RequestMetadata, RequestsContainer,
8    },
9    Result,
10};
11
12/// Builder for synchronous batch content generation requests
13pub struct BatchBuilder {
14    client: Arc<GeminiClient>,
15    display_name: String,
16    requests: Vec<GenerateContentRequest>,
17}
18
19impl BatchBuilder {
20    /// Create a new batch builder
21    pub(crate) fn new(client: Arc<GeminiClient>) -> Self {
22        Self {
23            client,
24            display_name: "RustBatch".to_string(),
25            requests: Vec::new(),
26        }
27    }
28
29    /// Set the display name for the batch request
30    pub fn with_name(mut self, name: String) -> Self {
31        self.display_name = name;
32        self
33    }
34
35    /// Set the requests for the batch request
36    pub fn with_requests(mut self, requests: Vec<GenerateContentRequest>) -> Self {
37        self.requests = requests;
38        self
39    }
40
41    /// Add a request to the batch request
42    pub fn with_request(mut self, request: GenerateContentRequest) -> Self {
43        self.requests.push(request);
44        self
45    }
46
47    /// Build the batch request
48    pub fn build(self) -> BatchGenerateContentRequest {
49        let batch_requests: Vec<BatchRequestItem> = self
50            .requests
51            .into_iter()
52            .enumerate()
53            .map(|(i, request)| BatchRequestItem {
54                request,
55                metadata: Some(RequestMetadata {
56                    key: format!("request-{}", i + 1),
57                }),
58            })
59            .collect();
60
61        BatchGenerateContentRequest {
62            batch: BatchConfig {
63                display_name: self.display_name,
64                input_config: InputConfig {
65                    requests: RequestsContainer {
66                        requests: batch_requests,
67                    },
68                },
69            },
70        }
71    }
72
73    /// Execute the batch request
74    pub async fn execute(self) -> Result<BatchGenerateContentResponse> {
75        let client = self.client.clone();
76        let request = self.build();
77        client.batch_generate_content_sync(request).await
78    }
79}
gemini_rust/batch_builder.rs

gemini_rust/
batch_builder.rs
