From 9b28ea89fd6072d04b6dc05a343aea7c665af9cd Mon Sep 17 00:00:00 2001 From: Tyler Longwell Date: Thu, 23 Oct 2025 12:38:32 -0400 Subject: [PATCH] Remove ModelOverrideProvider decorator in favor of Agent-based model config threading --- crates/goose-cli/src/session/builder.rs | 42 +++- crates/goose/src/agents/agent.rs | 28 ++- crates/goose/src/agents/reply_parts.rs | 37 ++- .../agents/subagent_execution_tool/tasks.rs | 25 +-- crates/goose/src/agents/subagent_handler.rs | 19 +- crates/goose/src/providers/mod.rs | 2 - crates/goose/src/providers/model_override.rs | 210 ------------------ 7 files changed, 116 insertions(+), 247 deletions(-) delete mode 100644 crates/goose/src/providers/model_override.rs diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs index f67fb8904d..badaf0d28e 100644 --- a/crates/goose-cli/src/session/builder.rs +++ b/crates/goose-cli/src/session/builder.rs @@ -224,12 +224,32 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { let temperature = session_config.settings.as_ref().and_then(|s| s.temperature); - let model_config = goose::model::ModelConfig::new(&model_name) - .unwrap_or_else(|e| { + // Check if we're doing a model-only override (no provider change) + let is_model_only_override = session_config.settings.as_ref().map_or(false, |s| { + s.goose_provider.is_none() && (s.goose_model.is_some() || s.temperature.is_some()) + }); + + // For model-only overrides, create provider with default model, then apply override to agent + // For provider changes, create provider with the new model config + let provider_model_config = if is_model_only_override { + // Get the default model for this provider + let default_model = config + .get_param::("GOOSE_MODEL") + .ok() + .expect("No model configured. Run 'goose configure' first"); + goose::model::ModelConfig::new(&default_model).unwrap_or_else(|e| { output::render_error(&format!("Failed to create model configuration: {}", e)); process::exit(1); }) - .with_temperature(temperature); + } else { + // Use the full model config including overrides + goose::model::ModelConfig::new(&model_name) + .unwrap_or_else(|e| { + output::render_error(&format!("Failed to create model configuration: {}", e)); + process::exit(1); + }) + .with_temperature(temperature) + }; // Create the agent let agent: Agent = Agent::new(); @@ -242,7 +262,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { ) .await; - let new_provider = match create(&provider_name, model_config).await { + let new_provider = match create(&provider_name, provider_model_config).await { Ok(provider) => provider, Err(e) => { output::render_error(&format!( @@ -278,6 +298,20 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { process::exit(1); }); + // Apply model-only override if needed + if is_model_only_override { + let override_config = goose::model::ModelConfig::new(&model_name) + .unwrap_or_else(|e| { + output::render_error(&format!( + "Failed to create override model configuration: {}", + e + )); + process::exit(1); + }) + .with_temperature(temperature); + agent.set_model_override(override_config).await; + } + // Handle session resolution and resuming let session_id: Option = if session_config.no_session { None diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 13460a646d..8487c6e22b 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -60,6 +60,7 @@ use super::platform_tools; use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE}; use crate::agents::subagent_task_config::TaskConfig; use crate::conversation::message::{Message, MessageContent, SystemNotificationType, ToolRequest}; +use crate::model::ModelConfig; use crate::session::extension_data::{EnabledExtensionsState, ExtensionState}; use crate::session::SessionManager; @@ -87,6 +88,7 @@ pub struct ToolCategorizeResult { /// The main goose Agent pub struct Agent { pub(super) provider: Mutex>>, + pub(super) model_config_override: Mutex>, pub extension_manager: Arc, pub(super) sub_recipe_manager: Mutex, pub(super) tasks_manager: TasksManager, @@ -162,6 +164,7 @@ impl Agent { Self { provider: Mutex::new(None), + model_config_override: Mutex::new(None), extension_manager: Arc::new(ExtensionManager::new()), sub_recipe_manager: Mutex::new(SubRecipeManager::new()), tasks_manager: TasksManager::new(), @@ -360,6 +363,29 @@ impl Agent { } } + /// Set a model config override for this agent + pub async fn set_model_override(&self, model_config: ModelConfig) { + let mut override_lock = self.model_config_override.lock().await; + *override_lock = Some(model_config); + } + + /// Clear any model config override + pub async fn clear_model_override(&self) { + let mut override_lock = self.model_config_override.lock().await; + *override_lock = None; + } + + /// Get the effective model config (override or provider's default) + pub async fn effective_model_config(&self) -> Result { + let override_lock = self.model_config_override.lock().await; + if let Some(ref override_config) = *override_lock { + Ok(override_config.clone()) + } else { + let provider = self.provider().await?; + Ok(provider.get_model_config()) + } + } + /// Check if a tool is a frontend tool pub async fn is_frontend_tool(&self, name: &str) -> bool { self.frontend_tools.lock().await.contains_key(name) @@ -973,7 +999,7 @@ impl Agent { } } - let mut stream = Self::stream_response_from_provider( + let mut stream = self.stream_response_from_provider( self.provider().await?, &system_prompt, conversation.messages(), diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index bc443d8374..1a3839c40e 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -101,13 +101,18 @@ impl Agent { /// Stream a response from the LLM provider. /// Handles toolshim transformations if needed pub(crate) async fn stream_response_from_provider( + &self, provider: Arc, system_prompt: &str, messages: &[Message], tools: &[Tool], toolshim_tools: &[Tool], ) -> Result { - let config = provider.get_model_config(); + // Check if we have a model override + let model_config_override = self.model_config_override.lock().await.clone(); + let config = model_config_override + .clone() + .unwrap_or_else(|| provider.get_model_config()); // Convert tool messages to text if toolshim is enabled let messages_for_provider = if config.toolshim { @@ -124,7 +129,8 @@ impl Agent { // Capture errors during stream creation and return them as part of the stream // so they can be handled by the existing error handling logic in the agent - let stream_result = if provider.supports_streaming() { + let stream_result = if provider.supports_streaming() && model_config_override.is_none() { + // Only stream if no override (for safety, matching ModelOverrideProvider behavior) debug!("WAITING_LLM_STREAM_START"); let result = provider .stream( @@ -137,13 +143,26 @@ impl Agent { result } else { debug!("WAITING_LLM_START"); - let complete_result = provider - .complete( - system_prompt.as_str(), - messages_for_provider.messages(), - &tools, - ) - .await; + let complete_result = if let Some(ref override_config) = model_config_override { + // Use complete_with_model when we have an override + provider + .complete_with_model( + override_config, + system_prompt.as_str(), + messages_for_provider.messages(), + &tools, + ) + .await + } else { + // Use regular complete when no override + provider + .complete( + system_prompt.as_str(), + messages_for_provider.messages(), + &tools, + ) + .await + }; debug!("WAITING_LLM_END"); match complete_result { diff --git a/crates/goose/src/agents/subagent_execution_tool/tasks.rs b/crates/goose/src/agents/subagent_execution_tool/tasks.rs index 9583b281ec..9a34888577 100644 --- a/crates/goose/src/agents/subagent_execution_tool/tasks.rs +++ b/crates/goose/src/agents/subagent_execution_tool/tasks.rs @@ -50,12 +50,7 @@ async fn handle_recipe_task( // Apply recipe provider settings if specified if let Some(ref settings) = recipe.settings { - // Validate: provider requires model - if settings.goose_provider.is_some() && settings.goose_model.is_none() { - return Err("Recipe specifies provider but no model".to_string()); - } - - // Apply settings based on what's specified + // Handle full provider replacement if let (Some(provider), Some(model)) = (&settings.goose_provider, &settings.goose_model) { // Full replacement: new provider and model let model_config = @@ -63,21 +58,11 @@ async fn handle_recipe_task( task_config.provider = providers::create(provider, model_config) .await .map_err(|e| format!("Failed to create provider '{}': {}", provider, e))?; - } else if settings.goose_model.is_some() || settings.temperature.is_some() { - // Partial override: wrap existing provider - let mut model_config = task_config.provider.get_model_config(); - if let Some(model) = &settings.goose_model { - model_config.model_name = model.clone(); - } - if let Some(temp) = settings.temperature { - model_config = model_config.with_temperature(Some(temp)); - } - task_config.provider = Arc::new(providers::ModelOverrideProvider::new( - task_config.provider.clone(), - model_config, - )); + } else if settings.goose_provider.is_some() { + // Provider without model is invalid + return Err("Recipe specifies provider but no model".to_string()); } - // else: no changes needed, keep current provider + // Note: Model-only overrides (without provider change) are handled in subagent_handler.rs via set_model_override() } let result = tokio::select! { diff --git a/crates/goose/src/agents/subagent_handler.rs b/crates/goose/src/agents/subagent_handler.rs index a7aca1236a..3754dbc3ae 100644 --- a/crates/goose/src/agents/subagent_handler.rs +++ b/crates/goose/src/agents/subagent_handler.rs @@ -122,10 +122,27 @@ fn get_agent_messages(recipe: Recipe, task_config: TaskConfig) -> AgentMessagesF .await .map_err(|e| anyhow!("Failed to get sub agent session file path: {}", e))?; agent - .update_provider(task_config.provider) + .update_provider(task_config.provider.clone()) .await .map_err(|e| anyhow!("Failed to set provider on sub agent: {}", e))?; + // Apply model override from recipe settings if present + // This handles model-only overrides (when provider is not changed) + if let Some(ref settings) = recipe.settings { + if settings.goose_provider.is_none() + && (settings.goose_model.is_some() || settings.temperature.is_some()) + { + let mut model_config = task_config.provider.get_model_config(); + if let Some(model) = &settings.goose_model { + model_config.model_name = model.clone(); + } + if let Some(temp) = settings.temperature { + model_config = model_config.with_temperature(Some(temp)); + } + agent.set_model_override(model_config).await; + } + } + for extension in task_config.extensions { if let Err(e) = agent.add_extension(extension.clone()).await { debug!( diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 1890f2faaf..d50502a42b 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -18,7 +18,6 @@ pub mod githubcopilot; pub mod google; pub mod lead_worker; pub mod litellm; -pub mod model_override; pub mod oauth; pub mod ollama; pub mod openai; @@ -38,4 +37,3 @@ pub mod venice; pub mod xai; pub use factory::{create, create_with_named_model, providers, refresh_custom_providers}; -pub use model_override::ModelOverrideProvider; diff --git a/crates/goose/src/providers/model_override.rs b/crates/goose/src/providers/model_override.rs deleted file mode 100644 index 4b97fe335a..0000000000 --- a/crates/goose/src/providers/model_override.rs +++ /dev/null @@ -1,210 +0,0 @@ -use std::sync::Arc; - -use async_trait::async_trait; - -use super::base::{ - LeadWorkerProviderTrait, MessageStream, Provider, ProviderMetadata, ProviderUsage, -}; -use super::errors::ProviderError; -use super::retry::RetryConfig; -use crate::conversation::message::Message; -use crate::model::ModelConfig; -use rmcp::model::Tool; - -/// A provider wrapper that overrides the model configuration -pub struct ModelOverrideProvider { - inner: Arc, - model_config: ModelConfig, -} - -impl ModelOverrideProvider { - pub fn new(inner: Arc, model_config: ModelConfig) -> Self { - Self { - inner, - model_config, - } - } -} - -#[async_trait] -impl Provider for ModelOverrideProvider { - fn metadata() -> ProviderMetadata - where - Self: Sized, - { - // This won't be called for instances - unimplemented!("ModelOverrideProvider doesn't have static metadata") - } - - fn get_model_config(&self) -> ModelConfig { - self.model_config.clone() - } - - async fn complete_with_model( - &self, - model_config: &ModelConfig, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - // Use the provided config (for complete_fast, etc.) - self.inner - .complete_with_model(model_config, system, messages, tools) - .await - } - - async fn complete( - &self, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - // Use our override config - self.inner - .complete_with_model(&self.model_config, system, messages, tools) - .await - } - - // Streaming: Disable for model override cases - fn supports_streaming(&self) -> bool { - false // Force non-streaming path for overrides - } - - async fn stream( - &self, - _system: &str, - _messages: &[Message], - _tools: &[Tool], - ) -> Result { - Err(ProviderError::NotImplemented( - "Streaming not supported with model override".to_string(), - )) - } - - // Delegate all other methods to inner provider - async fn complete_fast( - &self, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - self.inner.complete_fast(system, messages, tools).await - } - - fn retry_config(&self) -> RetryConfig { - self.inner.retry_config() - } - - async fn fetch_supported_models(&self) -> Result>, ProviderError> { - self.inner.fetch_supported_models().await - } - - fn supports_embeddings(&self) -> bool { - self.inner.supports_embeddings() - } - - async fn create_embeddings(&self, texts: Vec) -> Result>, ProviderError> { - self.inner.create_embeddings(texts).await - } - - fn as_lead_worker(&self) -> Option<&dyn LeadWorkerProviderTrait> { - None // Override providers are never lead-worker - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::conversation::message::Message; - use crate::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage}; - use crate::providers::errors::ProviderError; - use async_trait::async_trait; - use rmcp::model::Tool; - - // Simple mock provider for testing - struct MockProvider { - model_config: ModelConfig, - } - - impl MockProvider { - fn new(model_name: String) -> Self { - Self { - model_config: ModelConfig::new_or_fail(&model_name), - } - } - } - - #[async_trait] - impl Provider for MockProvider { - fn metadata() -> ProviderMetadata - where - Self: Sized, - { - ProviderMetadata::new( - "mock", - "Mock Provider", - "Test provider", - "mock-model", - vec![], - "https://example.com", - vec![], - ) - } - - fn get_model_config(&self) -> ModelConfig { - self.model_config.clone() - } - - async fn complete_with_model( - &self, - _model_config: &ModelConfig, - _system: &str, - _messages: &[Message], - _tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - Ok(( - Message::assistant().with_text("Mock response"), - ProviderUsage::new("mock-model".to_string(), Usage::default()), - )) - } - } - - #[tokio::test] - async fn test_model_override_provider() { - // Create a test provider with a default model - let inner_provider = Arc::new(MockProvider::new("gpt-4".to_string())); - let original_config = inner_provider.get_model_config(); - assert_eq!(original_config.model_name, "gpt-4"); - - // Create override with different model - let mut override_config = original_config.clone(); - override_config.model_name = "gpt-3.5-turbo".to_string(); - override_config = override_config.with_temperature(Some(0.7)); - - let override_provider = ModelOverrideProvider::new(inner_provider.clone(), override_config); - - // Verify the override config is returned - let config = override_provider.get_model_config(); - assert_eq!(config.model_name, "gpt-3.5-turbo"); - assert_eq!(config.temperature, Some(0.7)); - - // Verify streaming is disabled - assert!(!override_provider.supports_streaming()); - } - - #[tokio::test] - async fn test_complete_uses_override() { - let inner_provider = Arc::new(MockProvider::new("gpt-4".to_string())); - - let mut override_config = inner_provider.get_model_config(); - override_config.model_name = "claude-3".to_string(); - - let override_provider = ModelOverrideProvider::new(inner_provider, override_config.clone()); - - // When calling complete(), it should use the override config - let result = override_provider.complete("test", &[], &[]).await; - - assert!(result.is_ok()); - // The MockProvider doesn't actually use the model, but we've verified the flow - } -}