diff --git a/crates/goose-provider-types/src/formats.rs b/crates/goose-provider-types/src/formats.rs index 3db82a6ea..03f759096 100644 --- a/crates/goose-provider-types/src/formats.rs +++ b/crates/goose-provider-types/src/formats.rs @@ -1,5 +1,6 @@ pub mod anthropic; pub mod databricks; +pub mod google; pub mod ollama; pub mod openai; pub mod openai_responses; diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose-provider-types/src/formats/google.rs similarity index 96% rename from crates/goose/src/providers/formats/google.rs rename to crates/goose-provider-types/src/formats/google.rs index b8a5e2770..8d9d5e913 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose-provider-types/src/formats/google.rs @@ -1,9 +1,9 @@ +use crate::conversation::token_usage::{ProviderUsage, Usage}; +use crate::errors::ProviderError; +use crate::formats::openai::{is_valid_function_name, sanitize_function_name}; +use crate::model::ModelConfig; +use crate::thinking::ThinkingEffort; use anyhow::Result; -use goose_providers::conversation::token_usage::{ProviderUsage, Usage}; -use goose_providers::errors::ProviderError; -use goose_providers::formats::openai::{is_valid_function_name, sanitize_function_name}; -use goose_providers::model::ModelConfig; -use goose_providers::thinking::ThinkingEffort; use rmcp::model::{ object, AnnotateAble, CallToolRequestParams, ErrorCode, ErrorData, RawContent, Role, Tool, }; @@ -17,7 +17,7 @@ use std::ops::Deref; pub const THOUGHT_SIGNATURE_KEY: &str = "thoughtSignature"; const SYNTHETIC_THOUGHT_SIGNATURE: &str = "skip_thought_signature_validator"; -const GEMINI25_DEFAULT_THINKING_BUDGET: i32 = 8192; +const DEFAULT_THINKING_BUDGET: i32 = 8192; pub fn metadata_with_signature(signature: &str) -> ProviderMetadata { let mut map = ProviderMetadata::new(); @@ -537,7 +537,10 @@ struct GoogleRequest<'a> { generation_config: Option, } -fn get_thinking_config(model_config: &ModelConfig) -> Option { +fn get_thinking_config( + model_config: &ModelConfig, + thinking_budget: Option, +) -> Option { if model_config.reasoning == Some(false) || model_config.thinking_effort() == Some(ThinkingEffort::Off) { @@ -585,22 +588,19 @@ fn get_thinking_config(model_config: &ModelConfig) -> Option { } else { let thinking_budget = match model_config .request_param::("thinking_budget") - .or_else(|| { - crate::config::Config::global() - .get_param("GEMINI25_THINKING_BUDGET") - .ok() - }) { + .or(thinking_budget) + { Some(budget) if budget >= 0 => budget, Some(budget) => { tracing::warn!( "Invalid thinking budget '{}' for model '{}'. Must be >= 0. Using '{}'.", budget, model_config.model_name, - GEMINI25_DEFAULT_THINKING_BUDGET, + DEFAULT_THINKING_BUDGET, ); - GEMINI25_DEFAULT_THINKING_BUDGET + DEFAULT_THINKING_BUDGET } - None => GEMINI25_DEFAULT_THINKING_BUDGET, + None => DEFAULT_THINKING_BUDGET, }; Some(ThinkingConfig { thinking_level: None, @@ -615,6 +615,26 @@ pub fn create_request( system: &str, messages: &[Message], tools: &[Tool], +) -> Result { + create_request_impl(model_config, system, messages, tools, None) +} + +pub fn create_request_with_thinking_budget( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + thinking_budget: Option, +) -> Result { + create_request_impl(model_config, system, messages, tools, thinking_budget) +} + +fn create_request_impl( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + thinking_budget: Option, ) -> Result { let tools_wrapper = if tools.is_empty() { None @@ -624,7 +644,7 @@ pub fn create_request( }) }; - let thinking_config = get_thinking_config(model_config); + let thinking_config = get_thinking_config(model_config, thinking_budget); let generation_config = Some(GenerationConfig { temperature: model_config.temperature.map(|t| t as f64), @@ -1417,27 +1437,27 @@ data: [DONE]"#; #[test] fn test_get_thinking_config_disabled_reasoning() { - use goose_providers::model::ModelConfig; + use crate::model::ModelConfig; let config = ModelConfig::new("gemini-2.5-flash").with_thinking_effort(ThinkingEffort::Off); - let thinking_config = get_thinking_config(&config).unwrap(); + let thinking_config = get_thinking_config(&config, None).unwrap(); assert_eq!(thinking_config.thinking_budget, Some(0)); assert!(!thinking_config.include_thoughts); let config = ModelConfig::new("gemini-2.5-pro").with_thinking_effort(ThinkingEffort::Off); - assert!(get_thinking_config(&config).is_none()); + assert!(get_thinking_config(&config, None).is_none()); } #[test] fn test_get_thinking_config() { - use goose_providers::model::ModelConfig; + use crate::model::ModelConfig; // Test 1: Gemini 3 model with low thinking effort let mut params = std::collections::HashMap::new(); params.insert("thinking_effort".to_string(), serde_json::json!("low")); let mut config = ModelConfig::new("gemini-3-pro"); config.request_params = Some(params); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_some()); let thinking_config = result.unwrap(); assert!(thinking_config.thinking_level.is_some()); @@ -1449,7 +1469,7 @@ data: [DONE]"#; params.insert("thinking_effort".to_string(), serde_json::json!("high")); let mut config = ModelConfig::new("Gemini-3-Flash"); config.request_params = Some(params); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_some()); let thinking_config = result.unwrap(); assert!(matches!( @@ -1458,20 +1478,20 @@ data: [DONE]"#; )); let config = ModelConfig::new("gemini-2.5-flash"); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_some()); let thinking_config = result.unwrap(); assert!(thinking_config.include_thoughts); assert!(thinking_config.thinking_level.is_none()); assert_eq!( thinking_config.thinking_budget, - Some(GEMINI25_DEFAULT_THINKING_BUDGET) + Some(DEFAULT_THINKING_BUDGET) ); let mut params = HashMap::new(); params.insert("thinking_budget".to_string(), json!(4096)); let config = ModelConfig::new("gemini-2.5-flash").with_merged_request_params(params); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_some()); let thinking_config = result.unwrap(); assert_eq!(thinking_config.thinking_budget, Some(4096)); @@ -1479,20 +1499,20 @@ data: [DONE]"#; let mut params = HashMap::new(); params.insert("thinking_budget".to_string(), json!(-1)); let config = ModelConfig::new("gemini-2.5-flash").with_merged_request_params(params); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_some()); let thinking_config = result.unwrap(); assert_eq!( thinking_config.thinking_budget, - Some(GEMINI25_DEFAULT_THINKING_BUDGET) + Some(DEFAULT_THINKING_BUDGET) ); let config = ModelConfig::new("gemini-2.0-flash"); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_none()); let config = ModelConfig::new("gpt-4o"); - let result = get_thinking_config(&config); + let result = get_thinking_config(&config, None); assert!(result.is_none()); } } diff --git a/crates/goose/src/providers/google.rs b/crates/goose-providers/src/google.rs similarity index 79% rename from crates/goose/src/providers/google.rs rename to crates/goose-providers/src/google.rs index d00a40f0b..934a488cc 100644 --- a/crates/goose/src/providers/google.rs +++ b/crates/goose-providers/src/google.rs @@ -1,19 +1,18 @@ -use super::api_client::{ApiClient, AuthMethod}; -use super::base::MessageStream; -use super::openai_compatible::{handle_status, map_http_error_to_provider_error, sanitize_url}; -use super::retry::ProviderRetry; +use crate::api_client::{ApiClient, AuthMethod}; +use crate::base::MessageStream; use crate::conversation::message::Message; -use goose_providers::errors::ProviderError; +use crate::errors::ProviderError; +use crate::openai_compatible::{handle_status, map_http_error_to_provider_error, sanitize_url}; +use crate::retry::ProviderRetry; -use crate::providers::base::{ConfigKey, Provider, ProviderDef, ProviderMetadata}; -use crate::providers::formats::google::{create_request, response_to_streaming_message}; +use crate::base::{ConfigKey, Provider, ProviderMetadata}; +use crate::formats::google::{create_request_with_thinking_budget, response_to_streaming_message}; +use crate::model::ModelConfig; +use crate::request_log::{start_log, LoggerHandleExt}; use anyhow::Result; use async_stream::try_stream; use async_trait::async_trait; -use futures::future::BoxFuture; use futures::TryStreamExt; -use goose_providers::model::ModelConfig; -use goose_providers::request_log::{start_log, LoggerHandleExt}; use rmcp::model::Tool; use serde_json::Value; use std::io; @@ -22,7 +21,7 @@ use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LinesCodec}; use tokio_util::io::StreamReader; -pub(crate) const GOOGLE_PROVIDER_NAME: &str = "google"; +pub const GOOGLE_PROVIDER_NAME: &str = "google"; pub const GOOGLE_API_HOST: &str = "https://generativelanguage.googleapis.com"; pub const GOOGLE_DEFAULT_MODEL: &str = "gemini-2.5-pro"; pub const GOOGLE_DEFAULT_FAST_MODEL: &str = "gemini-2.5-flash"; @@ -62,30 +61,33 @@ pub struct GoogleProvider { api_client: ApiClient, #[serde(skip)] name: String, + #[serde(skip)] + thinking_budget: Option, } impl GoogleProvider { - pub async fn from_env( - tls_config: Option, + pub fn new( + host: String, + api_key: String, + tls_config: Option, + request_builder: Option, + thinking_budget: Option, ) -> Result { - let config = crate::config::Config::global(); - let api_key: String = config.get_secret("GOOGLE_API_KEY")?; - let host: String = config - .get_param("GOOGLE_HOST") - .unwrap_or_else(|_| GOOGLE_API_HOST.to_string()); - let auth = AuthMethod::ApiKey { header_name: "x-goog-api-key".to_string(), key: api_key, }; - let api_client = ApiClient::new_with_tls(host, auth, tls_config)? - .with_request_builder(crate::session_context::session_id_request_builder()) + let mut api_client = ApiClient::new_with_tls(host, auth, tls_config)? .with_header("Content-Type", "application/json")?; + if let Some(request_builder) = request_builder { + api_client = api_client.with_request_builder(request_builder); + } Ok(Self { api_client, name: GOOGLE_PROVIDER_NAME.to_string(), + thinking_budget, }) } @@ -100,7 +102,7 @@ impl GoogleProvider { } } -impl goose_providers::base::ProviderDescriptor for GoogleProvider { +impl crate::base::ProviderDescriptor for GoogleProvider { fn metadata() -> ProviderMetadata { ProviderMetadata::new( GOOGLE_PROVIDER_NAME, @@ -124,17 +126,6 @@ impl goose_providers::base::ProviderDescriptor for GoogleProvider { } } -impl ProviderDef for GoogleProvider { - type Provider = Self; - - fn from_env( - _extensions: Vec, - tls_config: Option, - ) -> BoxFuture<'static, Result> { - Box::pin(Self::from_env(tls_config)) - } -} - #[async_trait] impl Provider for GoogleProvider { fn get_name(&self) -> &str { @@ -180,7 +171,13 @@ impl Provider for GoogleProvider { messages: &[Message], tools: &[Tool], ) -> Result { - let payload = create_request(model_config, system, messages, tools)?; + let payload = create_request_with_thinking_budget( + model_config, + system, + messages, + tools, + self.thinking_budget, + )?; let mut log = start_log(model_config, &payload)?; let response = self diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index 2c1e3a8e7..77e011b17 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -3,6 +3,7 @@ pub mod api_client; pub mod databricks; pub mod databricks_auth; pub mod databricks_v2; +pub mod google; pub use goose_provider_types::{ base, canonical, conversation, errors, formats, goose_mode, images, json, model, permission, request_log, retry, thinking, utils, diff --git a/crates/goose/src/providers/formats/mod.rs b/crates/goose/src/providers/formats/mod.rs index 9f5531670..8f86a9373 100644 --- a/crates/goose/src/providers/formats/mod.rs +++ b/crates/goose/src/providers/formats/mod.rs @@ -7,7 +7,28 @@ pub mod databricks { pub use goose_providers::formats::databricks::*; } pub mod gcpvertexai; -pub mod google; +pub mod google { + use anyhow::Result; + use goose_providers::conversation::message::Message; + pub use goose_providers::formats::google::*; + use goose_providers::model::ModelConfig; + use rmcp::model::Tool; + use serde_json::Value; + + use crate::config::Config; + + pub fn create_request( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + // TODO: Remove this config fallback wrapper once gemini_oauth and Vertex/GCP Gemini + // move into goose-providers and receive provider config during construction. + let thinking_budget = Config::global().get_param("GEMINI25_THINKING_BUDGET").ok(); + create_request_with_thinking_budget(model_config, system, messages, tools, thinking_budget) + } +} pub mod openrouter; pub mod snowflake { pub use goose_providers::formats::snowflake::*; diff --git a/crates/goose/src/providers/google_def.rs b/crates/goose/src/providers/google_def.rs new file mode 100644 index 000000000..5bd3d9814 --- /dev/null +++ b/crates/goose/src/providers/google_def.rs @@ -0,0 +1,45 @@ +use anyhow::Result; +use futures::future::BoxFuture; +use goose_providers::api_client::TlsConfig; +use goose_providers::base::{ProviderDescriptor, ProviderMetadata}; +use goose_providers::google::{GoogleProvider, GOOGLE_API_HOST}; + +use crate::config::{Config, ExtensionConfig}; +use crate::providers::base::ProviderDef; + +pub struct GoogleProviderDef; + +impl ProviderDescriptor for GoogleProviderDef { + fn metadata() -> ProviderMetadata { + GoogleProvider::metadata() + } +} + +impl ProviderDef for GoogleProviderDef { + type Provider = GoogleProvider; + + fn from_env( + _extensions: Vec, + tls_config: Option, + ) -> BoxFuture<'static, Result> { + Box::pin(from_env(tls_config)) + } +} + +pub async fn from_env(tls_config: Option) -> Result { + let config = Config::global(); + let api_key: String = config.get_secret("GOOGLE_API_KEY")?; + let host: String = config + .get_param("GOOGLE_HOST") + .unwrap_or_else(|_| GOOGLE_API_HOST.to_string()); + + let thinking_budget = config.get_param("GEMINI25_THINKING_BUDGET").ok(); + + GoogleProvider::new( + host, + api_key, + tls_config, + Some(crate::session_context::session_id_request_builder()), + thinking_budget, + ) +} diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index c0d80d364..3ce282f7c 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -23,7 +23,6 @@ use super::{ gemini_cli::GeminiCliProvider, gemini_oauth::GeminiOAuthProvider, githubcopilot::GithubCopilotProvider, - google::GoogleProvider, huggingface::HuggingFaceProvider, kimicode::KimiCodeProvider, litellm::LiteLLMProvider, @@ -41,6 +40,7 @@ use crate::providers::anthropic_def::AnthropicProviderDef; use crate::providers::base::ProviderType; use crate::providers::databricks_def::{self, DatabricksProviderDef}; use crate::providers::databricks_v2_def::{self, DatabricksV2ProviderDef}; +use crate::providers::google_def::GoogleProviderDef; use crate::providers::ollama_def::OllamaProviderDef; use crate::providers::openai_def::OpenAiProviderDef; use crate::{ @@ -104,7 +104,7 @@ async fn init_registry() -> RwLock { registry.register::(false); registry.register::(true); registry.register::(false); - registry.register_with_inventory::( + registry.register_with_inventory::( true, Some(registrations::google_inventory()), ); diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index c6ecf4561..c0a961b01 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -37,7 +37,10 @@ pub mod gcpvertexai; pub mod gemini_cli; pub mod gemini_oauth; pub mod githubcopilot; -pub mod google; +pub mod google { + pub use goose_providers::google::*; +} +pub mod google_def; pub mod http_status { pub use goose_providers::http_status::*; }