gemini_rust/
batch.rs

1//! The Batch module for managing batch operations.
2//!
3//! For more information, see the official Google AI documentation:
4//! - [Batch Mode Guide](https://ai.google.dev/gemini-api/docs/batch-mode)
5//! - [Batch API Reference](https://ai.google.dev/api/batch-mode)
6//!
7//! # Design Note: Resource Management in Batch Operations
8//!
9//! The Batch API methods that consume the [`Batch`] struct (`cancel`, `delete`, `wait_for_completion`)
10//! return `std::result::Result<T, (Self, crate::Error)>` instead of the crate's `Result<T>`.
11//! This design follows patterns used in channel libraries (e.g., `std::sync::mpsc::Receiver`)
12//! and provides two key benefits:
13//!
14//! 1. **Resource Safety**: Once a [`Batch`] is consumed by an operation, it cannot be used again,
15//!    preventing invalid operations on deleted or canceled batches.
16//!
17//! 2. **Error Recovery**: If an operation fails due to transient network issues, both the
18//!    [`Batch`] and error information are returned, allowing callers to retry the operation.
19//!
20//! ## Example usage:
21//! ```rust,ignore
22//! use gemini_rust::{Gemini, Message};
23//!
24//! #[tokio::main]
25//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
26//!     let client = Gemini::new(std::env::var("GEMINI_API_KEY")?);
27//!     let request = client.generate_content().with_user_message("Why is the sky blue?").build();
28//!     let batch = client.batch_generate_content_sync().with_request(request).execute().await?;
29//!
30//!     match batch.delete().await {
31//!         Ok(()) => println!("Batch deleted successfully!"),
32//!         Err((batch, error)) => {
33//!             println!("Failed to delete batch: {}", error);
34//!             // Can retry: batch.delete().await
35//!         }
36//!     }
37//!     Ok(())
38//! }
39//! ```
40
41use std::{sync::Arc, time::Duration};
42use tokio::time::sleep;
43
44use crate::{
45    client::GeminiClient,
46    models::{BatchOperation, BatchStatus},
47    Error, Result,
48};
49
50/// Represents a long-running batch operation, providing methods to manage its lifecycle.
51///
52/// A `Batch` object is a handle to a batch operation on the Gemini API. It allows you to
53/// check the status, cancel the operation, or delete it once it's no longer needed.
54pub struct Batch {
55    /// The unique resource name of the batch operation, e.g., `operations/batch-xxxxxxxx`.
56    pub name: String,
57    client: Arc<GeminiClient>,
58}
59
60impl Batch {
61    /// Creates a new Batch instance.
62    pub(crate) fn new(name: String, client: Arc<GeminiClient>) -> Self {
63        Self { name, client }
64    }
65
66    /// Returns the unique resource name of the batch operation.
67    pub fn name(&self) -> &str {
68        &self.name
69    }
70
71    /// Retrieves the current status of the batch operation by making an API call.
72    ///
73    /// This method provides a snapshot of the batch's state at a single point in time.
74    pub async fn status(&self) -> Result<BatchStatus> {
75        let operation: BatchOperation = self.client.get_batch_operation(&self.name).await?;
76        BatchStatus::from_operation(operation)
77    }
78
79    /// Sends a request to the API to cancel the batch operation.
80    ///
81    /// Cancellation is not guaranteed to be instantaneous. The operation may continue to run for
82    /// some time after the cancellation request is made.
83    ///
84    /// Consumes the batch. If cancellation fails, returns the batch and error information
85    /// so it can be retried.
86    pub async fn cancel(self) -> std::result::Result<(), (Self, crate::Error)> {
87        match self.client.cancel_batch_operation(&self.name).await {
88            Ok(()) => Ok(()),
89            Err(e) => Err((self, e)),
90        }
91    }
92
93    /// Deletes the batch operation resource from the server.
94    ///
95    /// Note: This method indicates that the client is no longer interested in the operation result.
96    /// It does not cancel a running operation. To stop a running batch, use the `cancel` method.
97    /// This method should typically be used after the batch has completed.
98    ///
99    /// Consumes the batch. If deletion fails, returns the batch and error information
100    /// so it can be retried.
101    pub async fn delete(self) -> std::result::Result<(), (Self, crate::Error)> {
102        match self.client.delete_batch_operation(&self.name).await {
103            Ok(()) => Ok(()),
104            Err(e) => Err((self, e)),
105        }
106    }
107
108    /// Waits for the batch operation to complete by periodically polling its status.
109    ///
110    /// This method polls the batch status with a specified delay until the operation
111    /// reaches a terminal state (Succeeded, Failed, Cancelled, or Expired).
112    ///
113    /// Consumes the batch and returns the final status. If there's an error during polling,
114    /// the batch is returned in the error variant so it can be retried.
115    pub async fn wait_for_completion(
116        self,
117        delay: Duration,
118    ) -> std::result::Result<BatchStatus, (Self, crate::Error)> {
119        let batch_name = self.name.clone();
120        loop {
121            match self.status().await {
122                Ok(status) => match status {
123                    BatchStatus::Succeeded { .. } | BatchStatus::Cancelled => return Ok(status),
124                    BatchStatus::Expired => {
125                        return Err((self, Error::BatchExpired { name: batch_name }))
126                    }
127                    _ => sleep(delay).await,
128                },
129                Err(e) => match e {
130                    Error::BatchFailed { .. } => return Err((self, e)),
131                    _ => return Err((self, e)), // Return the batch and error for retry
132                },
133            }
134        }
135    }
136}
gemini_rust/batch.rs

gemini_rust/
batch.rs
