google_gemini_rs/
lib.rs

1pub mod client;
2pub mod google;
3
4pub use rust_mcp_sdk;
5
6#[cfg(test)]
7mod test {
8    use dotenv::dotenv;
9    use std::{env, path::Path};
10    use thiserror::Error;
11
12    use crate::{
13        client::{self, Client},
14        google::{self, common::Modality},
15    };
16
17    const GEMINI_API_ENV_KEY: &str = "GEMINI_API_KEY";
18    const GEMINI_MODEL_ENV_KEY: &str = "GEMINI_MODEL";
19
20    const TUX_IMAGE_PATH: &str = "tests/images/tux.png";
21
22    #[derive(Debug, Error)]
23    enum Error {
24        #[error(transparent)]
25        DotEnv(#[from] dotenv::Error),
26        #[error(transparent)]
27        Client(#[from] client::Error),
28        #[error(transparent)]
29        Var(#[from] env::VarError),
30        #[error(transparent)]
31        Io(#[from] std::io::Error),
32        #[error(transparent)]
33        Google(#[from] google::Error),
34    }
35
36    async fn client() -> Result<Client, Error> {
37        dotenv()?;
38
39        let key = env::var(GEMINI_API_ENV_KEY)?;
40        let model = env::var(GEMINI_MODEL_ENV_KEY)?;
41
42        Ok(Client::new(&model.as_str().try_into()?, &key)
43            .await?
44            .with_defaults())
45    }
46
47    #[tokio::test]
48    async fn basic_query() -> Result<(), Error> {
49        let mut client = client().await?;
50        let response = client
51            .send_text("I'm new to AI, so introduce yourself.")
52            .await?;
53        println!("{}", response.text().expect("Expected text output."));
54        let response = client.send_text("Hello.").await?;
55        println!("{}", response.text().expect("Expected text output."));
56        let response = client.send_text("Do you have a name?").await?;
57        println!("{}", response.text().expect("Expected text output."));
58        Ok(())
59    }
60
61    #[tokio::test]
62    async fn image_query() -> Result<(), Error> {
63        let mut client = client().await?;
64
65        let response = client
66            .send_text("Generate a thumbnail sized picture of a capybara.")
67            .await?;
68        println!("Image response: {:?}", response.images());
69        if client.model.output.contains(&Modality::Image) {
70            response
71                .images()
72                .first()
73                .expect("Expected image output(s).");
74        }
75        println!("Text response: {:?}", response.text());
76        Ok(())
77    }
78
79    #[tokio::test]
80    async fn image_and_text_query() -> Result<(), Error> {
81        let mut client = client().await?;
82        let response = client
83            .send_text("Your role is an artists that upgrades logos.")
84            .await?;
85
86        println!("{:?}", response.text().expect("Expected text result."));
87
88        let pic = Path::new(TUX_IMAGE_PATH);
89        let response = client.send_image_file(Some("Here is an image of the linux mascot, tux.  Add the words linux to the background".to_string()), &pic).await?;
90
91        println!("Response text: {:?}", response.text());
92
93        if client.model.output.contains(&Modality::Image) {
94            println!(
95                "{:?}",
96                response.images().first().expect("Expected image output(s)")
97            );
98        }
99
100        let response = client
101            .send_text("What type of animal is in the image you sent?")
102            .await?;
103
104        let text = response.text().expect("Expected text result.");
105
106        println!("Response check text: {}", text);
107
108        assert!(text.contains("penguin"));
109
110        Ok(())
111    }
112}
google_gemini_rs/lib.rs

google_gemini_rs/
lib.rs
