tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

move google provider into goose-providers (#10216)

Browse Source
This commit is contained in:
Jack Amadeo
2026-07-05 18:51:36 -07:00
committed by GitHub
parent 4e2a992c5a
commit 8e18f07487
8 changed files with 155 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-provider-types/src/formats.rs
+1
View File
@@ -1,5 +1,6 @@
pub mod anthropic;
pub mod databricks;
pub mod google;
pub mod ollama;
pub mod openai;
pub mod openai_responses;
crates/goose/src/providers/formats/google.rs → crates/goose-provider-types/src/formats/google.rs
+49 -29
View File
@@ -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<GenerationConfig>,
}
fn get_thinking_config(model_config: &ModelConfig) -> Option<ThinkingConfig> {
fn get_thinking_config(
model_config: &ModelConfig,
thinking_budget: Option<i32>,
) -> Option<ThinkingConfig> {
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<ThinkingConfig> {
} else {
let thinking_budget = match model_config
.request_param::<i32>("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<Value> {
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<i32>,
) -> Result<Value> {
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<i32>,
) -> Result<Value> {
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());
}
}
crates/goose/src/providers/google.rs → crates/goose-providers/src/google.rs
+31 -34
View File
@@ -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<i32>,
}
impl GoogleProvider {
pub async fn from_env(
tls_config: Option<crate::providers::api_client::TlsConfig>,
pub fn new(
host: String,
api_key: String,
tls_config: Option<crate::api_client::TlsConfig>,
request_builder: Option<crate::api_client::RequestBuilderDecorator>,
thinking_budget: Option<i32>,
) -> Result<Self> {
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<crate::config::ExtensionConfig>,
tls_config: Option<crate::providers::api_client::TlsConfig>,
) -> BoxFuture<'static, Result<Self::Provider>> {
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<MessageStream, ProviderError> {
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
crates/goose-providers/src/lib.rs
+1
View File
@@ -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,
crates/goose/src/providers/formats/mod.rs
+22 -1
View File
@@ -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<Value> {
// 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::*;
crates/goose/src/providers/google_def.rs
+45
View File
@@ -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<ExtensionConfig>,
tls_config: Option<TlsConfig>,
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(from_env(tls_config))
}
}
pub async fn from_env(tls_config: Option<TlsConfig>) -> Result<GoogleProvider> {
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,
)
}
crates/goose/src/providers/init.rs
+2 -2
View File
@@ -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<ProviderRegistry> {
registry.register::<GeminiCliProvider>(false);
registry.register::<GeminiOAuthProvider>(true);
registry.register::<GithubCopilotProvider>(false);
registry.register_with_inventory::<GoogleProvider>(
registry.register_with_inventory::<GoogleProviderDef>(
true,
Some(registrations::google_inventory()),
);
crates/goose/src/providers/mod.rs
+4 -1
View File
@@ -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::*;
}