awesome-bookmarks/goose
3
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-17 12:56:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove ModelOverrideProvider decorator in favor of Agent-based model config threading

Browse Source
This commit is contained in:
Tyler Longwell
2025-10-23 12:38:32 -04:00
parent dde49b5616
commit 9b28ea89fd
7 changed files with 116 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/builder.rs
+38 -4
View File
@@ -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::<String>("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<String> = if session_config.no_session {
None
crates/goose/src/agents/agent.rs
+27 -1
View File
@@ -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<Option<Arc<dyn Provider>>>,
pub(super) model_config_override: Mutex<Option<ModelConfig>>,
pub extension_manager: Arc<ExtensionManager>,
pub(super) sub_recipe_manager: Mutex<SubRecipeManager>,
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<ModelConfig> {
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(),
crates/goose/src/agents/reply_parts.rs
+28 -9
View File
@@ -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<dyn Provider>,
system_prompt: &str,
messages: &[Message],
tools: &[Tool],
toolshim_tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
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 {
crates/goose/src/agents/subagent_execution_tool/tasks.rs
+5 -20
View File
@@ -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! {
crates/goose/src/agents/subagent_handler.rs
+18 -1
View File
@@ -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!(
crates/goose/src/providers/mod.rs
-2
View File
@@ -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;
crates/goose/src/providers/model_override.rs
-210
View File
@@ -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<dyn Provider>,
model_config: ModelConfig,
}
impl ModelOverrideProvider {
pub fn new(inner: Arc<dyn Provider>, 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<MessageStream, ProviderError> {
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<Option<Vec<String>>, ProviderError> {
self.inner.fetch_supported_models().await
}
fn supports_embeddings(&self) -> bool {
self.inner.supports_embeddings()
}
async fn create_embeddings(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>, 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
}
}