use text::text_generator::{TextGenerator, TextGeneratorImpl};
mod config;
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 text::request::{Content, Part, RequestBody, SafetySetting};
use super::*;
#[tokio::test]
async fn test_api_call() {
let api_key = config::Config::from_env().unwrap().api_key.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 = config::Config::from_env().unwrap().api_key.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_interactice_chat() {
let api_key = config::Config::from_env().unwrap().api_key.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());
}
}