Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: adding oauth for gemini #1007

Closed
wants to merge 5 commits into from
Closed

feat: adding oauth for gemini #1007

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8adfbcd
feat: adding oauth for gemini
damienrj Jan 31, 2025
7516cdf
Add logging to log file says when are where oauth token is used
damienrj Feb 3, 2025
d87a608
Add more clear error messages
damienrj Feb 3, 2025
3228570
set api as required for configure
damienrj Feb 3, 2025
4e08422
Update crates/goose/src/providers/google.rs
damienrj Feb 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/goose/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ serde_yaml = "0.9.34"
once_cell = "1.20.2"
dirs = "6.0.0"
rand = "0.8.5"
oauth2 = { version = "4.4", features = ["reqwest"] }
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not super familiar with this library, we use https://crates.io/crates/yup-oauth2 in the goose-mcp crate, do you know what distinctions there are?

some larger questions around the oauth flow:

  1. whats the expected duration of the token?
  2. if short-lived will the changes automatically re-attempt an oauth flow if tokens are expired?
  3. does/should this oauth2 crate and implementation handle refresh tokens?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can see if it will also work with yup-oauth2, if that is already an existing library.

As to the flow, the token lasts an hour I believe. If the token is expired it will get a new token. Does the databrick token also expire?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks to be about the same for databricks, about an hour and then we do the same oauth local browser flow

Yeah we use yup-oauth2 but in a slightly different context, for google's Desktop client OAuth flow, but hopefully it does work


[dev-dependencies]
criterion = "0.5"
Expand Down
139 changes: 123 additions & 16 deletions crates/goose/src/providers/google.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use super::errors::ProviderError;
use super::oauth::{self, DEFAULT_REDIRECT_URL};
use crate::message::Message;
use crate::model::ModelConfig;
use crate::providers::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
Expand All @@ -8,11 +9,15 @@ use anyhow::Result;
use async_trait::async_trait;
use mcp_core::tool::Tool;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;
use url::Url;

pub const GOOGLE_API_HOST: &str = "https://generativelanguage.googleapis.com";
pub const GOOGLE_AUTH_ENDPOINT: &str = "https://accounts.google.com/o/oauth2/v2/auth";
pub const GOOGLE_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
const DEFAULT_SCOPES: &[&str] = &["https://www.googleapis.com/auth/generative-language.retriever"];
pub const GOOGLE_DEFAULT_MODEL: &str = "gemini-2.0-flash-exp";
pub const GOOGLE_KNOWN_MODELS: &[&str] = &[
"models/gemini-1.5-pro-latest",
Expand All @@ -25,12 +30,38 @@ pub const GOOGLE_KNOWN_MODELS: &[&str] = &[

pub const GOOGLE_DOC_URL: &str = "https://ai.google/get-started/our-models/";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GoogleAuth {
ApiKey(String),
OAuth {
client_id: String,
client_secret: String,
redirect_url: String,
scopes: Vec<String>,
},
}

impl GoogleAuth {
pub fn api_key(key: String) -> Self {
Self::ApiKey(key)
}

pub fn oauth(client_id: String, client_secret: String) -> Self {
Self::OAuth {
client_id,
client_secret,
redirect_url: DEFAULT_REDIRECT_URL.to_string(),
scopes: DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect(),
}
}
}

#[derive(Debug, serde::Serialize)]
pub struct GoogleProvider {
#[serde(skip)]
client: Client,
host: String,
api_key: String,
auth: GoogleAuth,
model: ModelConfig,
}

Expand All @@ -44,7 +75,6 @@ impl Default for GoogleProvider {
impl GoogleProvider {
pub fn from_env(model: ModelConfig) -> Result<Self> {
let config = crate::config::Config::global();
let api_key: String = config.get_secret("GOOGLE_API_KEY")?;
let host: String = config
.get("GOOGLE_HOST")
.unwrap_or_else(|_| GOOGLE_API_HOST.to_string());
Expand All @@ -53,12 +83,74 @@ impl GoogleProvider {
.timeout(Duration::from_secs(600))
.build()?;

Ok(Self {
client,
host,
api_key,
model,
})
// First try API key authentication
if let Ok(api_key) = config.get_secret("GOOGLE_API_KEY") {
return Ok(Self {
client,
host,
auth: GoogleAuth::api_key(api_key),
model,
});
}

// Try OAuth authentication
let client_id = config.get("GOOGLE_CLIENT_ID");
let client_secret = config.get_secret("GOOGLE_CLIENT_SECRET");

match (client_id, client_secret) {
(Ok(id), Ok(secret)) => {
let scopes = DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect();
Ok(Self {
client,
host,
auth: GoogleAuth::OAuth {
client_id: id,
client_secret: secret,
redirect_url: DEFAULT_REDIRECT_URL.to_string(),
scopes,
},
model,
})
}
_ => Err(anyhow::anyhow!(
"Authentication configuration missing. Please set either GOOGLE_API_KEY or both GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET"
)),
}
}

async fn ensure_auth_header(&self) -> Result<String, ProviderError> {
match &self.auth {
GoogleAuth::ApiKey(key) => Ok(format!("Bearer {}", key)),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to format this into a Bearer {} type token if it is only used in the url params key={api_key} ?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can give that a try, went through a few iterations before I got it working.

GoogleAuth::OAuth {
client_id,
client_secret,
scopes,
.. // Ignore redirect_url as we're using the default
} => {
let token = if client_secret.is_empty() {
// Use public client OAuth if no client secret
oauth::get_oauth_token_public_client_async(
GOOGLE_AUTH_ENDPOINT,
GOOGLE_TOKEN_ENDPOINT,
client_id,
scopes,
).await
} else {
// Use private client OAuth if client secret is present
oauth::get_oauth_token_with_endpoints_async(
GOOGLE_AUTH_ENDPOINT,
GOOGLE_TOKEN_ENDPOINT,
client_id,
client_secret,
scopes,
).await
};

token
.map_err(|e| ProviderError::Authentication(format!("Failed to get OAuth token: {}", e)))
.map(|token| format!("Bearer {}", token))
}
}
}

async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
Expand All @@ -67,20 +159,33 @@ impl GoogleProvider {

let url = base_url
.join(&format!(
"v1beta/models/{}:generateContent?key={}",
self.model.model_name, self.api_key
"v1beta/models/{}:generateContent",
self.model.model_name,
))
.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;

let response = self
let auth = self.ensure_auth_header().await?;

// Add auth either as query param for API key or header for OAuth
let mut request = self
.client
.post(url)
.header("CONTENT_TYPE", "application/json")
.json(&payload)
.send()
.await?;
.post(url.to_string())
.header("Content-Type", "application/json");

match &self.auth {
GoogleAuth::ApiKey(_) => {
// Remove "Bearer " prefix for API key and pass as query param
let api_key = auth.trim_start_matches("Bearer ").to_string();
request = request.query(&[("key", api_key)]);
}
GoogleAuth::OAuth { .. } => {
request = request.header("Authorization", auth);
}
}

let response = request.json(&payload).send().await?;

let status = response.status();
let payload: Option<Value> = response.json().await.ok();
Expand Down Expand Up @@ -136,6 +241,8 @@ impl Provider for GoogleProvider {
vec![
ConfigKey::new("GOOGLE_API_KEY", true, true, None),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How would the user be able to select oauth as an authentication method if the api key is required first?

image

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently via environment variable, some design choices need to be made for how configure/UI work if a Provider has more than one option. So while that is determined left configure and UI to only use the API token.

ConfigKey::new("GOOGLE_HOST", false, false, Some(GOOGLE_API_HOST)),
ConfigKey::new("GOOGLE_CLIENT_ID", false, false, None),
ConfigKey::new("GOOGLE_CLIENT_SECRET", false, true, None),
],
)
}
Expand Down
Loading
Loading