text_request_stream/
text_request_stream.rs

1use std::env;
2use std::io::{stdout, Write};
3
4use google_generative_ai_rs::v1::{
5    api::Client,
6    gemini::{request::Request, response::GeminiResponse, Content, Part, ResponseType, Role},
7};
8
9/// Simple text request using the public API and an API key for authn
10/// To run:
11/// ```
12/// API_KEY=[YOUR_API_KEY] RUST_LOG=info cargo run --package google-generative-ai-rs  --example text_request
13/// ``
14#[tokio::main]
15async fn main() -> Result<(), Box<dyn std::error::Error>> {
16    env_logger::init();
17
18    let token = match env::var("API_KEY") {
19        Ok(v) => v,
20        Err(e) => {
21            let msg = "$API_KEY not found".to_string();
22            panic!("{e:?}:{msg}");
23        }
24    };
25
26    // Either run as a standard text request or a stream generate content request
27    let client = Client::new_from_model_response_type(
28        google_generative_ai_rs::v1::gemini::Model::Gemini1_0Pro,
29        token.clone(),
30        ResponseType::StreamGenerateContent,
31    );
32
33    println!("token {:#?}", token);
34
35    let txt_request = Request {
36        contents: vec![Content {
37            role: Role::User,
38            parts: vec![Part {
39                text: Some("Give me a recipe for banana bread.".to_string()),
40                inline_data: None,
41                file_data: None,
42                video_metadata: None,
43            }],
44        }],
45        tools: vec![],
46        safety_settings: vec![],
47        generation_config: None,
48
49        #[cfg(feature = "beta")]
50        system_instruction: None,
51    };
52
53    let response = client.post(30, &txt_request).await?;
54
55    println!("output streaming content");
56
57    if let Some(stream_response) = response.streamed() {
58        if let Some(json_stream) = stream_response.response_stream {
59            Client::for_each_async(json_stream, move |response: GeminiResponse| async move {
60                let mut lock = stdout().lock();
61                write!(
62                    lock,
63                    "{}",
64                    response.candidates[0].content.parts[0]
65                        .text
66                        .clone()
67                        .unwrap()
68                        .as_str()
69                )
70                .unwrap();
71            })
72            .await
73        }
74    }
75
76    Ok(())
77}