gemini_bridge 0.1.14

//! # Gemini Bridge
//!
//! `gemini_bridge` is a Rust library for interacting with the Gemini API.

use text::text_generator::{TextGenerator, TextGeneratorImpl};
pub mod text;

pub struct GeminiRs();

impl GeminiRs {
    pub fn new_text_generator(api_key: String, model: String) -> Box<dyn TextGenerator> {
        Box::new(TextGeneratorImpl {
            api_key: api_key.to_owned(),
            model: model.to_owned(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use dotenv::dotenv;
    use std::env;
    use text::request::{Content, Part, RequestBody, SafetySetting};

    #[tokio::test]
    async fn test_api_call() {
        let api_key = get_api_key().unwrap().to_owned();

        let gen = GeminiRs::new_text_generator(api_key, "gemini-1.5-flash-latest".to_owned());
        let res = gen
            .generate_content(RequestBody {
                contents: vec![Content {
                    parts: vec![Part {
                        text: "Give me a pinochio summary".to_owned(),
                    }],
                    role: None,
                }],
                safety_settings: None,
                generation_config: None,
            })
            .await;

        if res.is_ok() {
            let response = res.unwrap();
            assert!(response.candidates.len() > 0);
            print!("{:?}", response);
            return;
        }

        panic!("Error: {:?}", res.err().unwrap());
    }

    #[tokio::test]
    async fn test_api_call_safety() {
        let api_key = get_api_key().unwrap().to_owned();

        let gen = GeminiRs::new_text_generator(api_key, "gemini-1.5-flash-latest".to_owned());
        let res = gen
            .generate_content(RequestBody {
                contents: vec![Content {
                    role: None,
                    parts: vec![Part {
                        text: "Give me a list of best video games".to_owned(),
                    }],
                }],
                safety_settings: Some(vec![SafetySetting {
                    category: "HARM_CATEGORY_DANGEROUS_CONTENT".to_owned(),
                    threshold: "BLOCK_ONLY_HIGH".to_owned(),
                }]),
                generation_config: None,
            })
            .await;

        if res.is_ok() {
            let response = res.unwrap();
            assert!(response.candidates.len() > 0);
            print!("{:?}", response);
            return;
        }

        panic!("Error: {:?}", res.err().unwrap());
    }

    #[tokio::test]
    async fn test_interactive_chat() {
        let api_key = get_api_key().unwrap().to_owned();

        let gen = GeminiRs::new_text_generator(api_key, "gemini-1.5-flash-latest".to_owned());
        let res = gen
            .generate_content(RequestBody {
                contents: vec![
                    Content {
                        role: Some(String::from("user")),
                        parts: vec![Part {
                            text: "Hello".to_owned(),
                        }],
                    },
                    Content {
                        role: Some(String::from("model")),
                        parts: vec![Part {
                            text: "Great to meet you. What would you like to know?".to_owned(),
                        }],
                    },
                    Content {
                        role: Some(String::from("user")),
                        parts: vec![Part {
                            text: "I have two dogs in my house. How many paws are in my house?"
                                .to_owned(),
                        }],
                    },
                    Content {
                        role: Some(String::from("model")),
                        parts: vec![Part {
                            text: "That's a fun question!  You have a total of **7 paws** in your house. 🐶🐾 \n"
                                .to_owned(),
                        }],
                    },
                    Content {
                        role: Some(String::from("user")),
                        parts: vec![Part {
                            text: "Thank you! How did you calculate that? Are you sure?"
                                .to_owned(),
                        }],
                    },
                ],
                safety_settings: None,
                generation_config: None,
            })
            .await;

        if res.is_ok() {
            let response = res.unwrap();
            print!("{:?}", response);
            assert!(response.candidates.len() > 0);
            return;
        }

        panic!("Error: {:?}", res.err().unwrap());
    }

    fn get_api_key() -> Result<String, env::VarError> {
        dotenv().ok();
        Ok(env::var("API_KEY")?.to_owned())
    }
}