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