gemini_rust/batch/
model.rs

1use serde::{Deserialize, Serialize};
2use snafu::Snafu;
3use time::OffsetDateTime;
4
5use crate::common::serde::*;
6use crate::generation::{GenerateContentRequest, GenerationResponse};
7use crate::Model;
8
9/// Batch file request line JSON representation.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct BatchRequestFileItem {
12    /// Batch generation request (wrapped in request field for API compatibility)
13    pub request: GenerateContentRequest,
14    /// Batch request unique identifier
15    #[serde(with = "key_as_string")]
16    pub key: usize,
17}
18
19/// Batch file response line JSON representation.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct BatchResponseFileItem {
22    /// Batch response (wrapped in response field for API compatibility)
23    #[serde(flatten)]
24    pub response: BatchGenerateContentResponseItem,
25    /// Batch response unique identifier
26    #[serde(with = "key_as_string")]
27    pub key: usize,
28}
29
30impl From<BatchGenerateContentResponseItem> for Result<GenerationResponse, IndividualRequestError> {
31    fn from(response: BatchGenerateContentResponseItem) -> Self {
32        match response {
33            BatchGenerateContentResponseItem::Response(r) => Ok(r),
34            BatchGenerateContentResponseItem::Error(err) => Err(err),
35        }
36    }
37}
38
39/// Represents the response of a batch operation.
40#[derive(Debug, Clone, Serialize, Deserialize)]
41#[serde(untagged)]
42pub enum BatchOperationResponse {
43    /// Response with inlined responses
44    #[serde(rename_all = "camelCase")]
45    InlinedResponses { inlined_responses: InlinedResponses },
46    /// Response with a file containing results
47    #[serde(rename_all = "camelCase")]
48    ResponsesFile { responses_file: String },
49}
50
51/// A container for inlined responses.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct InlinedResponses {
55    /// The list of batch response items
56    pub inlined_responses: Vec<InlinedBatchGenerationResponseItem>,
57}
58
59/// Represents a single response item within an inlined batch response.
60///
61/// This structure combines request metadata with the actual response or error,
62/// used when batch results are returned inline rather than in a separate file.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(rename_all = "camelCase")]
65pub struct InlinedBatchGenerationResponseItem {
66    /// Request metadata containing the original key and other identifiers
67    pub metadata: RequestMetadata,
68    /// The actual response content or error for this batch item
69    #[serde(flatten)]
70    pub result: BatchGenerateContentResponseItem,
71}
72
73/// An item in a batch generate content response.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75#[serde(rename_all = "camelCase")]
76pub enum BatchGenerateContentResponseItem {
77    /// Successful response item
78    Response(GenerationResponse),
79    /// Error response item
80    Error(IndividualRequestError),
81}
82
83/// An error for an individual request in a batch.
84#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
85pub struct IndividualRequestError {
86    /// The error code
87    pub code: i32,
88    /// The error message
89    pub message: String,
90    /// Additional details about the error
91    #[serde(skip_serializing_if = "Option::is_none")]
92    pub details: Option<serde_json::Value>,
93}
94
95/// Response from the Gemini API for batch content generation (async batch creation)
96#[derive(Debug, Clone, Serialize, Deserialize)]
97#[serde(rename_all = "camelCase")]
98pub struct BatchGenerateContentResponse {
99    /// The name/ID of the created batch
100    pub name: String,
101    /// Metadata about the batch
102    pub metadata: BatchMetadata,
103}
104
105/// Metadata for the batch operation
106#[derive(Debug, Clone, Serialize, Deserialize)]
107#[serde(rename_all = "camelCase")]
108pub struct BatchMetadata {
109    /// Type annotation
110    #[serde(rename = "@type")]
111    pub type_annotation: String,
112    /// Model used for the batch
113    pub model: Model,
114    /// Display name of the batch
115    pub display_name: String,
116    /// Creation time
117    #[serde(with = "time::serde::rfc3339")]
118    pub create_time: OffsetDateTime,
119    /// Update time
120    #[serde(with = "time::serde::rfc3339")]
121    pub update_time: OffsetDateTime,
122    /// Batch statistics
123    pub batch_stats: BatchStats,
124    /// Current state of the batch
125    pub state: BatchState,
126    /// Name of the batch (duplicate)
127    pub name: String,
128}
129
130/// Individual batch request item
131#[derive(Debug, Clone, Serialize, Deserialize)]
132#[serde(rename_all = "camelCase")]
133pub struct BatchRequestItem {
134    /// The actual request
135    pub request: GenerateContentRequest,
136    /// Metadata for the request
137    pub metadata: RequestMetadata,
138}
139
140/// Request for batch content generation (corrected format)
141#[derive(Debug, Clone, Serialize, Deserialize)]
142#[serde(rename_all = "camelCase")]
143pub struct BatchGenerateContentRequest {
144    /// The batch configuration
145    pub batch: BatchConfig,
146}
147
148/// Batch configuration
149#[derive(Debug, Clone, Serialize, Deserialize)]
150#[serde(rename_all = "camelCase")]
151pub struct BatchConfig {
152    /// Display name for the batch
153    pub display_name: String,
154    /// Input configuration
155    pub input_config: InputConfig,
156}
157
158/// The state of a batch operation.
159#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
160#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
161#[allow(clippy::enum_variant_names)]
162pub enum BatchState {
163    /// State is unspecified
164    BatchStateUnspecified,
165    /// Batch is pending execution
166    BatchStatePending,
167    /// Batch is currently running
168    BatchStateRunning,
169    /// Batch completed successfully
170    BatchStateSucceeded,
171    /// Batch failed during execution
172    BatchStateFailed,
173    /// Batch was cancelled
174    BatchStateCancelled,
175    /// Batch expired before completion
176    BatchStateExpired,
177}
178
179/// Statistics for the batch
180#[derive(Debug, Clone, Serialize, Deserialize)]
181#[serde(rename_all = "camelCase")]
182pub struct BatchStats {
183    /// Total number of requests in the batch
184    #[serde(with = "i64_as_string")]
185    pub request_count: i64,
186    /// Number of pending requests
187    #[serde(default, with = "i64_as_string::optional")]
188    pub pending_request_count: Option<i64>,
189    /// Number of completed requests
190    #[serde(default, with = "i64_as_string::optional")]
191    pub completed_request_count: Option<i64>,
192    /// Number of failed requests
193    #[serde(default, with = "i64_as_string::optional")]
194    pub failed_request_count: Option<i64>,
195    /// Number of successful requests
196    #[serde(default, with = "i64_as_string::optional")]
197    pub successful_request_count: Option<i64>,
198}
199
200/// Represents a long-running operation from the Gemini API.
201#[derive(Debug, Serialize, Deserialize)]
202#[serde(rename_all = "camelCase")]
203pub struct BatchOperation {
204    /// The resource name of the operation
205    pub name: String,
206    /// Metadata about the batch operation
207    pub metadata: BatchMetadata,
208    /// Whether the operation is complete
209    #[serde(default)]
210    pub done: bool,
211    /// The result of the operation (if complete)
212    #[serde(flatten)]
213    pub result: Option<OperationResult>,
214}
215
216/// Represents an error within a long-running operation.
217#[derive(Debug, Snafu, serde::Deserialize, serde::Serialize)]
218pub struct OperationError {
219    /// The error code
220    pub code: i32,
221    /// The error message
222    pub message: String,
223    // details are not included as they are not consistently typed in the API
224}
225
226/// Represents the result of a completed batch operation, which is either a response or an error.
227#[derive(Debug, Serialize, Deserialize)]
228#[serde(rename_all = "camelCase")]
229pub enum OperationResult {
230    /// Successful operation result
231    Response(BatchOperationResponse),
232    /// Failed operation result
233    Error(OperationError),
234}
235
236impl From<OperationResult> for Result<BatchOperationResponse, OperationError> {
237    fn from(operation: OperationResult) -> Self {
238        match operation {
239            OperationResult::Response(response) => Ok(response),
240            OperationResult::Error(error) => Err(error),
241        }
242    }
243}
244
245/// The outcome of a single request in a batch operation.
246/// Response from the Gemini API for listing batch operations.
247#[derive(Debug, serde::Deserialize)]
248#[serde(rename_all = "camelCase")]
249pub struct ListBatchesResponse {
250    /// A list of batch operations.
251    pub operations: Vec<BatchOperation>,
252    /// A token to retrieve the next page of results.
253    pub next_page_token: Option<String>,
254}
255
256/// Input configuration for batch requests
257#[derive(Debug, Clone, Serialize, Deserialize)]
258#[serde(rename_all = "camelCase")]
259pub enum InputConfig {
260    /// The requests to be processed in the batch.
261    Requests(RequestsContainer),
262    /// The name of the File containing the input requests.
263    FileName(String),
264}
265
266impl InputConfig {
267    /// Returns the batch size of the input configuration.
268    ///
269    /// Returns `None` if the input configuration is a file name.
270    pub fn batch_size(&self) -> Option<usize> {
271        match self {
272            InputConfig::Requests(container) => Some(container.requests.len()),
273            InputConfig::FileName(_) => None,
274        }
275    }
276}
277
278/// Container for requests
279#[derive(Debug, Clone, Serialize, Deserialize)]
280#[serde(rename_all = "camelCase")]
281pub struct RequestsContainer {
282    /// List of requests
283    pub requests: Vec<BatchRequestItem>,
284}
285
286/// Metadata for batch request
287#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
288#[serde(rename_all = "camelCase")]
289pub struct RequestMetadata {
290    /// Key for the request
291    #[serde(with = "key_as_string")]
292    pub key: usize,
293}