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