Skip to main content

streaming/
streaming.rs

1use futures::StreamExt;
2use gemini_client_api::gemini::ask::Gemini;
3use gemini_client_api::gemini::types::sessions::Session;
4use std::env;
5use std::io::{Write, stdout};
6
7#[tokio::main]
8async fn main() {
9    let mut session = Session::new(10);
10    let api_key = env::var("GEMINI_API_KEY").expect("GEMINI_API_KEY must be set");
11    let ai = Gemini::new(api_key, "gemini-2.5-flash", None);
12
13    println!("--- Streaming Example ---");
14    let prompt = "Write a poem about crab-like robots on Mars.";
15    println!("User: {}\n", prompt);
16    print!("Gemini: ");
17    stdout().flush().unwrap();
18    session.ask(prompt);
19
20    // Start a streaming request
21    let mut response_stream = ai.ask_as_stream(session).await.unwrap();
22
23    while let Some(chunk_result) = response_stream.next().await {
24        match chunk_result {
25            Ok(response) => {
26                // Get the text from the current chunk
27                let text = response.get_chat().get_text_no_think("");
28                print!("{text}");
29                stdout().flush().unwrap();
30            }
31            Err(e) => {
32                eprintln!("\nError receiving chunk: {e:?}",);
33                break;
34            }
35        }
36    }
37
38    println!("\n\n--- Stream Complete ---");
39    // Note: The session passed to ask_as_stream is updated as you exhaust the stream.
40    session = response_stream.get_session_owned();
41    println!("Updated session: {session:?}")
42}
43#[allow(dead_code)]
44async fn with_extractor() {
45    let mut session = Session::new(10);
46    let api_key = env::var("GEMINI_API_KEY").expect("GEMINI_API_KEY must be set");
47    let ai = Gemini::new(api_key, "gemini-2.5-flash", None);
48
49    println!("--- Streaming Example ---");
50    let prompt = "Write a poem about crab-like robots on Mars.";
51    println!("User: {prompt}\n");
52    println!("Gemini: ");
53    session.ask(prompt);
54
55    // Start a streaming request
56    let mut response_stream = ai
57        .ask_as_stream_with_extractor(session, |session, _gemini_response| {
58            session.get_last_chat().unwrap().get_text_no_think("\n")
59        })
60        .await
61        .unwrap();
62
63    while let Some(chunk_result) = response_stream.next().await {
64        match chunk_result {
65            Ok(response) => {
66                // Get the text from the current chunk
67                println!("Complete poem: {response}\n");
68            }
69            Err(e) => {
70                eprintln!("\nError receiving chunk: {e:?}",);
71                break;
72            }
73        }
74    }
75
76    println!("\n\n--- Stream Complete ---");
77    // Note: The session passed to ask_as_stream is updated as you exhaust the stream.
78    session = response_stream.get_session_owned();
79    println!("Updated session: {session:?}")
80}