feat: lead/worker model #2719

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.

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

Merged

michaelneale merged 24 commits into main from micn/lead-model

Jun 5, 2025

README.md

-Original file line number
+Diff line change
@@ Expand Up @@
     Designed for maximum flexibility, goose works with any LLM, seamlessly integrates with MCP servers, and is available as both a desktop app as well as CLI - making it the ultimate AI assistant for developers who want to move faster and focus on innovation.
+    ## Multiple Model Configuration
+    goose supports using different models for different purposes to optimize performance and cost, which can work across model providers as well as models.
+    ### Lead/Worker Model Pattern
+    Use a powerful model for initial planning and complex reasoning, then switch to a faster/cheaper model for execution, this happens automatically by goose:
+    ```bash
+    # Required: Enable lead model mode
+    export GOOSE_LEAD_MODEL=modelY
+    # Optional: configure a provider for the lead model if not the default provider
+    export GOOSE_LEAD_PROVIDER=providerX  # Defaults to main provider
+    ```
+    ### Planning Model Configuration
+    Use a specialized model for the `/plan` command in CLI mode, this is explicitly invoked when you want to plan (vs execute)
+    ```bash
+    # Optional: Use different model for planning
+    export GOOSE_PLANNER_PROVIDER=openai
+    export GOOSE_PLANNER_MODEL=gpt-4
+    ```
+    Both patterns help you balance model capabilities with cost and speed for optimal results, and switch between models and vendors as required.
     # Quick Links
     - [Quickstart](https://block.github.io/goose/docs/quickstart)
@@ Expand Down @@

crates/goose-cli/src/session/builder.rs

-Original file line number
+Diff line change
@@ Expand Up / @@ -7,6 +7,7 @@ use goose::session; @@
     use goose::session::Identifier;
     use mcp_client::transport::Error as McpClientError;
     use std::process;
+    use std::sync::Arc;
     use super::output;
     use super::Session;
@@ Expand Down Expand Up @@
         // Create the agent
         let agent: Agent = Agent::new();
         let new_provider = create(&provider_name, model_config).unwrap();
+        // Keep a reference to the provider for display_session_info
+        let provider_for_display = Arc::clone(&new_provider);
+        // Log model information at startup
+        if let Some(lead_worker) = new_provider.as_lead_worker() {
+            let (lead_model, worker_model) = lead_worker.get_model_info();
+            tracing::info!(
+                "🤖 Lead/Worker Mode Enabled: Lead model (first 3 turns): {}, Worker model (turn 4+): {}, Auto-fallback on failures: Enabled",
+                lead_model,
+                worker_model
+            );
+        } else {
+            tracing::info!("🤖 Using model: {}", model);
+        }
         agent
             .update_provider(new_provider)
             .await
@@ Expand Down Expand Up @@
             session.agent.override_system_prompt(override_prompt).await;
         }
-        output::display_session_info(session_config.resume, &provider_name, &model, &session_file);
+        output::display_session_info(
+            session_config.resume,
+            &provider_name,
+            &model,
+            &session_file,
+            Some(&provider_for_display),
+        );
         session
     }

crates/goose-cli/src/session/output.rs

-Original file line number
+Diff line change
@@ Expand Up / @@ -10,6 +10,7 @@ use std::cell::RefCell; @@
     use std::collections::HashMap;
     use std::io::Error;
     use std::path::Path;
+    use std::sync::Arc;
     use std::time::Duration;
     // Re-export theme for use in main
@@ Expand Down Expand Up / @@ -536,22 +537,56 @@ fn shorten_path(path: &str, debug: bool) -> String { @@
     }
     // Session display functions
-    pub fn display_session_info(resume: bool, provider: &str, model: &str, session_file: &Path) {
+    pub fn display_session_info(
+        resume: bool,
+        provider: &str,
+        model: &str,
+        session_file: &Path,
+        provider_instance: Option<&Arc<dyn goose::providers::base::Provider>>,
+    ) {
         let start_session_msg = if resume {
             "resuming session |"
         } else if session_file.to_str() == Some("/dev/null") || session_file.to_str() == Some("NUL") {
             "running without session |"
         } else {
             "starting session |"
         };
-        println!(
-            "{} {} {} {} {}",
-            style(start_session_msg).dim(),
-            style("provider:").dim(),
-            style(provider).cyan().dim(),
-            style("model:").dim(),
-            style(model).cyan().dim(),
-        );
+        // Check if we have lead/worker mode
+        if let Some(provider_inst) = provider_instance {
+            if let Some(lead_worker) = provider_inst.as_lead_worker() {
+                let (lead_model, worker_model) = lead_worker.get_model_info();
+                println!(
+                    "{} {} {} {} {} {} {}",
+                    style(start_session_msg).dim(),
+                    style("provider:").dim(),
+                    style(provider).cyan().dim(),
+                    style("lead model:").dim(),
+                    style(&lead_model).cyan().dim(),
+                    style("worker model:").dim(),
+                    style(&worker_model).cyan().dim(),
+                );
+            } else {
+                println!(
+                    "{} {} {} {} {}",
+                    style(start_session_msg).dim(),
+                    style("provider:").dim(),
+                    style(provider).cyan().dim(),
+                    style("model:").dim(),
+                    style(model).cyan().dim(),
+                );
+            }
+        } else {
+            // Fallback to original behavior if no provider instance
+            println!(
+                "{} {} {} {} {}",
+                style(start_session_msg).dim(),
+                style("provider:").dim(),
+                style(provider).cyan().dim(),
+                style("model:").dim(),
+                style(model).cyan().dim(),
+            );
+        }
         if session_file.to_str() != Some("/dev/null") && session_file.to_str() != Some("NUL") {
             println!(
@@ Expand Down @@

crates/goose/src/providers/base.rs

-Original file line number
+Diff line change
@@ Expand Up / @@ -148,6 +148,12 @@ impl Usage { @@
     use async_trait::async_trait;
+    /// Trait for LeadWorkerProvider-specific functionality
+    pub trait LeadWorkerProviderTrait {
+        /// Get information about the lead and worker models for logging
+        fn get_model_info(&self) -> (String, String);
+    }
     /// Base trait for AI providers (OpenAI, Anthropic, etc)
     #[async_trait]
     pub trait Provider: Send + Sync {
@@ Expand Down Expand Up / @@ -195,6 +201,12 @@ pub trait Provider: Send + Sync { @@
                 "This provider does not support embeddings".to_string(),
             ))
         }
+        /// Check if this provider is a LeadWorkerProvider
+        /// This is used for logging model information at startup
+        fn as_lead_worker(&self) -> Option<&dyn LeadWorkerProviderTrait> {
+            None
+        }
     }
     #[cfg(test)]
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: lead/worker model #2719

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!