batch_cancel/
batch_cancel.rs

1//! Example demonstrating how to cancel a batch operation when the user presses CTRL-C.
2//!
3//! This example shows:
4//! 1. Creating a batch operation with multiple requests
5//! 2. Setting up a signal handler for CTRL-C
6//! 3. Starting the batch operation
7//! 4. Canceling the batch when CTRL-C is pressed
8//! 5. Properly handling the result
9
10use gemini_rust::{Gemini, Message, Result};
11use std::{env, sync::Arc, time::Duration};
12use tokio::{signal, sync::Mutex};
13
14#[tokio::main]
15async fn main() -> Result<()> {
16    // Get the API key from the environment
17    let api_key = env::var("GEMINI_API_KEY").expect("GEMINI_API_KEY must be set");
18
19    // Create the Gemini client
20    let gemini = Gemini::new(api_key);
21
22    // Create a batch with multiple requests
23    let mut batch_generate_content = gemini
24        .batch_generate_content_sync()
25        .with_name("batch_cancel_example".to_string());
26
27    // Add several requests to make the batch take some time to process
28    for i in 1..=10 {
29        let request = gemini
30            .generate_content()
31            .with_message(Message::user(format!(
32                "Write a creative story about a robot learning to paint, part {}. Make it at least 100 words long.",
33                i
34            )))
35            .build();
36
37        batch_generate_content = batch_generate_content.with_request(request);
38    }
39
40    // Build and start the batch
41    let batch = batch_generate_content.execute().await?;
42    println!("Batch created successfully!");
43    println!("Batch Name: {}", batch.name());
44    println!("Press CTRL-C to cancel the batch operation...");
45
46    // Wrap the batch in an Arc<Mutex<Option<Batch>>> to allow safe sharing
47    let batch = Arc::new(Mutex::new(Some(batch)));
48    let batch_clone = Arc::clone(&batch);
49
50    // Spawn a task to handle CTRL-C
51    let cancel_task = tokio::spawn(async move {
52        // Wait for CTRL-C signal
53        signal::ctrl_c().await.expect("Failed to listen for CTRL-C");
54        println!("Received CTRL-C, canceling batch operation...");
55
56        // Take the batch from the Option, leaving None.
57        // The lock is released immediately after this block.
58        let mut batch_to_cancel = batch_clone.lock().await;
59
60        if let Some(batch) = batch_to_cancel.take() {
61            // Cancel the batch operation
62            match batch.cancel().await {
63                Ok(()) => {
64                    println!("Batch canceled successfully!");
65                }
66                Err((batch, e)) => {
67                    println!("Failed to cancel batch: {}. Retrying...", e);
68                    // Retry once
69                    match batch.cancel().await {
70                        Ok(()) => {
71                            println!("Batch canceled successfully on retry!");
72                        }
73                        Err((_, retry_error)) => {
74                            eprintln!("Failed to cancel batch even on retry: {}", retry_error);
75                        }
76                    }
77                }
78            }
79        } else {
80            println!("Batch was already processed.");
81        }
82    });
83
84    // Wait for a short moment to ensure the cancel task is ready
85    tokio::time::sleep(Duration::from_millis(100)).await;
86
87    // Wait for the batch to complete or be canceled
88    if let Some(batch) = batch.lock().await.take() {
89        println!("Waiting for batch to complete or be canceled...");
90        match batch.wait_for_completion(Duration::from_secs(5)).await {
91            Ok(final_status) => {
92                // Cancel task is no longer needed since batch completed
93                cancel_task.abort();
94
95                println!("Batch completed with status: {:?}", final_status);
96
97                // Print some details about the results
98                match final_status {
99                    gemini_rust::BatchStatus::Succeeded { .. } => {
100                        println!("Batch succeeded!");
101                    }
102                    gemini_rust::BatchStatus::Cancelled => {
103                        println!("Batch was cancelled as requested.");
104                    }
105                    gemini_rust::BatchStatus::Expired => {
106                        println!("Batch expired.");
107                    }
108                    _ => {
109                        println!("Batch finished with an unexpected status.");
110                    }
111                }
112            }
113            Err((batch, e)) => {
114                // This could happen if there was a network error while polling
115                println!("Error while waiting for batch completion: {}", e);
116
117                // Try one more time to get the status
118                match batch.status().await {
119                    Ok(status) => println!("Current batch status: {:?}", status),
120                    Err(status_error) => println!("Error getting final status: {}", status_error),
121                }
122            }
123        }
124    }
125
126    Ok(())
127}
batch_cancel/batch_cancel.rs

batch_cancel/
batch_cancel.rs
