diff --git a/examples/google-search-retrieval.rs b/examples/google-search-retrieval.rs new file mode 100644 index 0000000..275c15c --- /dev/null +++ b/examples/google-search-retrieval.rs @@ -0,0 +1,39 @@ +use gemini_rs::prelude::*; + +#[tokio::main] +async fn main() -> Result<(), Box> { + tracing_subscriber::fmt().init(); + let authentication_manager = gcp_auth::provider().await?; + let api_endpoint = std::env::var("API_ENDPOINT")?; + let project_id = std::env::var("PROJECT_ID")?; + let location_id = std::env::var("LOCATION_ID")?; + + let gemini = GeminiClient::new( + authentication_manager, + api_endpoint, + project_id, + location_id, + ); + + let prompt = "What day is today?"; + + let request = GenerateContentRequest { + contents: vec![Content { + role: Some("user".to_string()), + parts: Some(vec![Part::Text(prompt.to_string())]), + }], + tools: Some(vec![Tools { + google_search_retrieval: Some(GoogleSearchRetrieval::default()), + ..Default::default() + }]), + ..Default::default() + }; + + let result = gemini + .generate_content(&request, "gemini-1.0-pro-002") + .await?; + + println!("Response: {:?}", result.candidates[0].get_text().unwrap()); + + Ok(()) +} diff --git a/src/types/generate_content.rs b/src/types/generate_content.rs index 3b20214..9ad3efa 100644 --- a/src/types/generate_content.rs +++ b/src/types/generate_content.rs @@ -27,9 +27,17 @@ impl GenerateContentRequest { } } -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Default, Serialize, Deserialize)] pub struct Tools { pub function_declarations: Option>, + #[serde(rename = "googleSearchRetrieval")] + pub google_search_retrieval: Option, +} + +#[derive(Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct GoogleSearchRetrieval { + pub disable_attribution: bool, } #[derive(Clone, Debug, Serialize, Deserialize, Default)]