mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
Remove lead/worker provider (#7989)
Signed-off-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
@@ -648,18 +648,8 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
|
||||
process::exit(1);
|
||||
}
|
||||
};
|
||||
let provider_for_display = Arc::clone(&new_provider);
|
||||
|
||||
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: {}", resolved.model_name);
|
||||
}
|
||||
let provider_for_debug = Arc::clone(&new_provider);
|
||||
tracing::info!("🤖 Using model: {}", resolved.model_name);
|
||||
|
||||
agent
|
||||
.update_provider(new_provider, &session_id)
|
||||
@@ -692,7 +682,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
|
||||
let agent_ptr = resolve_and_load_extensions(
|
||||
agent,
|
||||
extensions_for_provider,
|
||||
Arc::clone(&provider_for_display),
|
||||
Arc::clone(&provider_for_debug),
|
||||
session_config.interactive,
|
||||
&session_id,
|
||||
)
|
||||
@@ -732,7 +722,6 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
|
||||
&resolved.provider_name,
|
||||
&resolved.model_name,
|
||||
&Some(session_id),
|
||||
Some(&provider_for_display),
|
||||
);
|
||||
}
|
||||
session
|
||||
|
||||
@@ -79,10 +79,6 @@ enum StreamEvent {
|
||||
#[serde(flatten)]
|
||||
data: NotificationData,
|
||||
},
|
||||
ModelChange {
|
||||
model: String,
|
||||
mode: String,
|
||||
},
|
||||
Error {
|
||||
error: String,
|
||||
},
|
||||
@@ -1070,13 +1066,6 @@ impl CliSession {
|
||||
Some(Ok(AgentEvent::HistoryReplaced(updated_conversation))) => {
|
||||
self.messages = updated_conversation;
|
||||
}
|
||||
Some(Ok(AgentEvent::ModelChange { model, mode })) => {
|
||||
if is_stream_json_mode {
|
||||
emit_stream_event(&StreamEvent::ModelChange { model: model.clone(), mode: mode.clone() });
|
||||
} else if self.debug {
|
||||
eprintln!("Model changed to {} in {} mode", model, mode);
|
||||
}
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
handle_agent_error(&e, is_stream_json_mode);
|
||||
cancel_token_clone.cancel();
|
||||
|
||||
@@ -17,7 +17,7 @@ use std::cell::RefCell;
|
||||
use std::collections::HashMap;
|
||||
use std::io::{Error, IsTerminal, Write};
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use std::sync::LazyLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::streaming_buffer::MarkdownBuffer;
|
||||
@@ -1251,7 +1251,6 @@ pub fn display_session_info(
|
||||
provider: &str,
|
||||
model: &str,
|
||||
session_id: &Option<String>,
|
||||
provider_instance: Option<&Arc<dyn goose::providers::base::Provider>>,
|
||||
) {
|
||||
set_terminal_title();
|
||||
|
||||
@@ -1263,16 +1262,7 @@ pub fn display_session_info(
|
||||
"new session"
|
||||
};
|
||||
|
||||
let model_display = 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();
|
||||
format!("{} → {}", lead_model, worker_model)
|
||||
} else {
|
||||
model.to_string()
|
||||
}
|
||||
} else {
|
||||
model.to_string()
|
||||
};
|
||||
let model_display = model.to_string();
|
||||
|
||||
let cwd_display = std::env::current_dir()
|
||||
.ok()
|
||||
|
||||
@@ -137,10 +137,6 @@ pub enum MessageEvent {
|
||||
reason: String,
|
||||
token_state: TokenState,
|
||||
},
|
||||
ModelChange {
|
||||
model: String,
|
||||
mode: String,
|
||||
},
|
||||
Notification {
|
||||
request_id: String,
|
||||
#[schema(value_type = Object)]
|
||||
@@ -364,9 +360,6 @@ pub async fn reply(
|
||||
stream_event(MessageEvent::UpdateConversation {conversation: new_messages}, &tx, &cancel_token).await;
|
||||
|
||||
}
|
||||
Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
|
||||
stream_event(MessageEvent::ModelChange { model, mode }, &tx, &cancel_token).await;
|
||||
}
|
||||
Ok(Some(Ok(AgentEvent::McpNotification((request_id, n))))) => {
|
||||
stream_event(MessageEvent::Notification{
|
||||
request_id: request_id.clone(),
|
||||
|
||||
@@ -462,13 +462,6 @@ pub async fn session_reply(
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Ok(Some(Ok(AgentEvent::ModelChange { model, mode }))) => {
|
||||
publish(
|
||||
Some(task_request_id.clone()),
|
||||
MessageEvent::ModelChange { model, mode },
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Ok(Some(Ok(AgentEvent::McpNotification((notification_request_id, n))))) => {
|
||||
publish(
|
||||
Some(task_request_id.clone()),
|
||||
|
||||
@@ -158,7 +158,6 @@ pub struct Agent {
|
||||
pub enum AgentEvent {
|
||||
Message(Message),
|
||||
McpNotification((String, ServerNotification)),
|
||||
ModelChange { model: String, mode: String },
|
||||
HistoryReplaced(Conversation),
|
||||
}
|
||||
|
||||
@@ -1230,27 +1229,6 @@ impl Agent {
|
||||
Ok((response, usage)) => {
|
||||
compaction_attempts = 0;
|
||||
|
||||
// Emit model change event if provider is lead-worker
|
||||
let provider = self.provider().await?;
|
||||
if let Some(lead_worker) = provider.as_lead_worker() {
|
||||
if let Some(ref usage) = usage {
|
||||
let active_model = usage.model.clone();
|
||||
let (lead_model, worker_model) = lead_worker.get_model_info();
|
||||
let mode = if active_model == lead_model {
|
||||
"lead"
|
||||
} else if active_model == worker_model {
|
||||
"worker"
|
||||
} else {
|
||||
"unknown"
|
||||
};
|
||||
|
||||
yield AgentEvent::ModelChange {
|
||||
model: active_model,
|
||||
mode: mode.to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref usage) = usage {
|
||||
self.update_session_metrics(&session_config.id, session_config.schedule_id.clone(), usage, false).await?;
|
||||
}
|
||||
|
||||
@@ -208,8 +208,6 @@ impl Agent {
|
||||
Ok((tools, toolshim_tools, system_prompt))
|
||||
}
|
||||
|
||||
// Don't add gen_ai.request.model here — provider.get_model_config()
|
||||
// returns the wrong model for LeadWorkerProvider.
|
||||
#[tracing::instrument(
|
||||
skip(provider, session_id, system_prompt, messages, tools, toolshim_tools),
|
||||
fields(session.id = %session_id)
|
||||
|
||||
@@ -208,7 +208,7 @@ fn get_agent_messages(params: SubagentRunParams) -> AgentMessagesFuture {
|
||||
}
|
||||
conversation.push(msg);
|
||||
}
|
||||
Ok(AgentEvent::McpNotification(_)) | Ok(AgentEvent::ModelChange { .. }) => {}
|
||||
Ok(AgentEvent::McpNotification(_)) => {}
|
||||
Ok(AgentEvent::HistoryReplaced(updated_conversation)) => {
|
||||
conversation = updated_conversation;
|
||||
}
|
||||
|
||||
@@ -445,17 +445,6 @@ impl GatewayHandler {
|
||||
"gateway stream: mcp notification #{event_count}"
|
||||
);
|
||||
}
|
||||
Ok(AgentEvent::ModelChange {
|
||||
ref model,
|
||||
ref mode,
|
||||
}) => {
|
||||
tracing::debug!(
|
||||
session_id,
|
||||
model,
|
||||
mode,
|
||||
"gateway stream: model change #{event_count}"
|
||||
);
|
||||
}
|
||||
Ok(AgentEvent::HistoryReplaced(_)) => {
|
||||
tracing::debug!(
|
||||
session_id,
|
||||
|
||||
@@ -421,30 +421,6 @@ async fn send_session_event(installation: &InstallationData) -> Result<(), Strin
|
||||
insert(&mut props, "setting_max_turns", max_turns);
|
||||
}
|
||||
|
||||
if let Ok(lead_model) = config.get_param::<String>("GOOSE_LEAD_MODEL") {
|
||||
insert(&mut props, "setting_lead_model", lead_model);
|
||||
}
|
||||
if let Ok(lead_provider) = config.get_param::<String>("GOOSE_LEAD_PROVIDER") {
|
||||
insert(&mut props, "setting_lead_provider", lead_provider);
|
||||
}
|
||||
if let Ok(lead_turns) = config.get_param::<i64>("GOOSE_LEAD_TURNS") {
|
||||
insert(&mut props, "setting_lead_turns", lead_turns);
|
||||
}
|
||||
if let Ok(lead_failure_threshold) = config.get_param::<i64>("GOOSE_LEAD_FAILURE_THRESHOLD") {
|
||||
insert(
|
||||
&mut props,
|
||||
"setting_lead_failure_threshold",
|
||||
lead_failure_threshold,
|
||||
);
|
||||
}
|
||||
if let Ok(lead_fallback_turns) = config.get_param::<i64>("GOOSE_LEAD_FALLBACK_TURNS") {
|
||||
insert(
|
||||
&mut props,
|
||||
"setting_lead_fallback_turns",
|
||||
lead_fallback_turns,
|
||||
);
|
||||
}
|
||||
|
||||
let extensions = get_enabled_extensions();
|
||||
insert(&mut props, "extensions_count", extensions.len() as u64);
|
||||
let extension_names: Vec<String> = extensions.iter().map(|e| e.name()).collect();
|
||||
|
||||
@@ -450,18 +450,6 @@ pub enum PermissionRouting {
|
||||
Noop,
|
||||
}
|
||||
|
||||
/// 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);
|
||||
|
||||
/// Get the currently active model name
|
||||
fn get_active_model(&self) -> String;
|
||||
|
||||
/// Get (lead_turns, failure_threshold, fallback_turns)
|
||||
fn get_settings(&self) -> (usize, usize, usize);
|
||||
}
|
||||
|
||||
/// Base trait for AI providers (OpenAI, Anthropic, etc)
|
||||
#[async_trait]
|
||||
pub trait Provider: Send + Sync {
|
||||
@@ -645,23 +633,6 @@ pub trait Provider: Send + Sync {
|
||||
))
|
||||
}
|
||||
|
||||
/// 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
|
||||
}
|
||||
|
||||
/// Get the currently active model name
|
||||
/// For regular providers, this returns the configured model
|
||||
/// For LeadWorkerProvider, this returns the currently active model (lead or worker)
|
||||
fn get_active_model_name(&self) -> String {
|
||||
if let Some(lead_worker) = self.as_lead_worker() {
|
||||
lead_worker.get_active_model()
|
||||
} else {
|
||||
self.get_model_config().model_name
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the first 3 user messages as strings for session naming
|
||||
fn get_initial_user_messages(&self, messages: &Conversation) -> Vec<String> {
|
||||
messages
|
||||
|
||||
@@ -17,7 +17,6 @@ use super::{
|
||||
gemini_cli::GeminiCliProvider,
|
||||
githubcopilot::GithubCopilotProvider,
|
||||
google::GoogleProvider,
|
||||
lead_worker::LeadWorkerProvider,
|
||||
litellm::LiteLLMProvider,
|
||||
local_inference::LocalInferenceProvider,
|
||||
nanogpt::NanoGptProvider,
|
||||
@@ -41,10 +40,6 @@ use crate::{
|
||||
use anyhow::Result;
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
const DEFAULT_LEAD_TURNS: usize = 3;
|
||||
const DEFAULT_FAILURE_THRESHOLD: usize = 2;
|
||||
const DEFAULT_FALLBACK_TURNS: usize = 2;
|
||||
|
||||
static REGISTRY: OnceCell<RwLock<ProviderRegistry>> = OnceCell::const_new();
|
||||
|
||||
async fn init_registry() -> RwLock<ProviderRegistry> {
|
||||
@@ -135,13 +130,6 @@ pub async fn create(
|
||||
model: ModelConfig,
|
||||
extensions: Vec<ExtensionConfig>,
|
||||
) -> Result<Arc<dyn Provider>> {
|
||||
let config = crate::config::Config::global();
|
||||
|
||||
if let Ok(lead_model_name) = config.get_param::<String>("GOOSE_LEAD_MODEL") {
|
||||
tracing::info!("Creating lead/worker provider from environment variables");
|
||||
return create_lead_worker_from_env(name, &model, &lead_model_name, extensions).await;
|
||||
}
|
||||
|
||||
let constructor = get_from_registry(name).await?.constructor.clone();
|
||||
constructor(model, extensions).await
|
||||
}
|
||||
@@ -179,179 +167,10 @@ pub async fn create_with_named_model(
|
||||
create(provider_name, config, extensions).await
|
||||
}
|
||||
|
||||
async fn create_lead_worker_from_env(
|
||||
default_provider_name: &str,
|
||||
default_model: &ModelConfig,
|
||||
lead_model_name: &str,
|
||||
extensions: Vec<ExtensionConfig>,
|
||||
) -> Result<Arc<dyn Provider>> {
|
||||
let config = crate::config::Config::global();
|
||||
|
||||
let lead_provider_name = config
|
||||
.get_param::<String>("GOOSE_LEAD_PROVIDER")
|
||||
.unwrap_or_else(|_| default_provider_name.to_string());
|
||||
|
||||
let lead_turns = config
|
||||
.get_param::<usize>("GOOSE_LEAD_TURNS")
|
||||
.unwrap_or(DEFAULT_LEAD_TURNS);
|
||||
let failure_threshold = config
|
||||
.get_param::<usize>("GOOSE_LEAD_FAILURE_THRESHOLD")
|
||||
.unwrap_or(DEFAULT_FAILURE_THRESHOLD);
|
||||
let fallback_turns = config
|
||||
.get_param::<usize>("GOOSE_LEAD_FALLBACK_TURNS")
|
||||
.unwrap_or(DEFAULT_FALLBACK_TURNS);
|
||||
|
||||
let lead_model_config = ModelConfig::new_with_context_env(
|
||||
lead_model_name.to_string(),
|
||||
&lead_provider_name,
|
||||
Some("GOOSE_LEAD_CONTEXT_LIMIT"),
|
||||
)?;
|
||||
|
||||
let worker_model_config = create_worker_model_config(default_model, default_provider_name)?;
|
||||
|
||||
let registry = get_registry().await;
|
||||
|
||||
let lead_constructor = {
|
||||
let guard = registry.read().unwrap();
|
||||
guard
|
||||
.entries
|
||||
.get(&lead_provider_name)
|
||||
.ok_or_else(|| anyhow::anyhow!("Unknown provider: {}", lead_provider_name))?
|
||||
.constructor
|
||||
.clone()
|
||||
};
|
||||
|
||||
let worker_constructor = {
|
||||
let guard = registry.read().unwrap();
|
||||
guard
|
||||
.entries
|
||||
.get(default_provider_name)
|
||||
.ok_or_else(|| anyhow::anyhow!("Unknown provider: {}", default_provider_name))?
|
||||
.constructor
|
||||
.clone()
|
||||
};
|
||||
|
||||
let lead_provider = lead_constructor(lead_model_config, extensions.clone()).await?;
|
||||
let worker_provider = worker_constructor(worker_model_config, extensions).await?;
|
||||
|
||||
Ok(Arc::new(LeadWorkerProvider::new_with_settings(
|
||||
lead_provider,
|
||||
worker_provider,
|
||||
lead_turns,
|
||||
failure_threshold,
|
||||
fallback_turns,
|
||||
)))
|
||||
}
|
||||
|
||||
fn create_worker_model_config(
|
||||
default_model: &ModelConfig,
|
||||
provider_name: &str,
|
||||
) -> Result<ModelConfig> {
|
||||
let mut worker_config = ModelConfig::new_or_fail(&default_model.model_name)
|
||||
.with_canonical_limits(provider_name)
|
||||
.with_context_limit(default_model.context_limit)
|
||||
.with_temperature(default_model.temperature)
|
||||
.with_max_tokens(default_model.max_tokens)
|
||||
.with_toolshim(default_model.toolshim)
|
||||
.with_toolshim_model(default_model.toolshim_model.clone());
|
||||
|
||||
let global_config = crate::config::Config::global();
|
||||
|
||||
if let Ok(limit) = global_config.get_param::<usize>("GOOSE_WORKER_CONTEXT_LIMIT") {
|
||||
worker_config = worker_config.with_context_limit(Some(limit));
|
||||
} else if let Ok(limit) = global_config.get_param::<usize>("GOOSE_CONTEXT_LIMIT") {
|
||||
worker_config = worker_config.with_context_limit(Some(limit));
|
||||
}
|
||||
|
||||
Ok(worker_config)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test_case::test_case(None, None, None, DEFAULT_LEAD_TURNS, DEFAULT_FAILURE_THRESHOLD, DEFAULT_FALLBACK_TURNS ; "defaults")]
|
||||
#[test_case::test_case(Some("7"), Some("4"), Some("3"), 7, 4, 3 ; "custom")]
|
||||
#[tokio::test]
|
||||
async fn test_create_lead_worker_provider(
|
||||
lead_turns: Option<&str>,
|
||||
failure_threshold: Option<&str>,
|
||||
fallback_turns: Option<&str>,
|
||||
expected_turns: usize,
|
||||
expected_failure: usize,
|
||||
expected_fallback: usize,
|
||||
) {
|
||||
let _guard = env_lock::lock_env([
|
||||
("GOOSE_LEAD_MODEL", Some("gpt-4o")),
|
||||
("GOOSE_LEAD_PROVIDER", None),
|
||||
("GOOSE_LEAD_TURNS", lead_turns),
|
||||
("GOOSE_LEAD_FAILURE_THRESHOLD", failure_threshold),
|
||||
("GOOSE_LEAD_FALLBACK_TURNS", fallback_turns),
|
||||
("OPENAI_API_KEY", Some("fake-openai-no-keyring")),
|
||||
("OPENAI_CUSTOM_HEADERS", Some("")),
|
||||
]);
|
||||
|
||||
let provider = create(
|
||||
"openai",
|
||||
ModelConfig::new_or_fail("gpt-4o-mini").with_canonical_limits("openai"),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let lw = provider.as_lead_worker().unwrap();
|
||||
let (lead, worker) = lw.get_model_info();
|
||||
assert_eq!(lead, "gpt-4o");
|
||||
assert_eq!(worker, "gpt-4o-mini");
|
||||
assert_eq!(
|
||||
lw.get_settings(),
|
||||
(expected_turns, expected_failure, expected_fallback)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_regular_provider_without_lead_config() {
|
||||
let _guard = env_lock::lock_env([
|
||||
("GOOSE_LEAD_MODEL", None),
|
||||
("GOOSE_LEAD_PROVIDER", None),
|
||||
("GOOSE_LEAD_TURNS", None),
|
||||
("GOOSE_LEAD_FAILURE_THRESHOLD", None),
|
||||
("GOOSE_LEAD_FALLBACK_TURNS", None),
|
||||
("OPENAI_API_KEY", Some("fake-openai-no-keyring")),
|
||||
("OPENAI_CUSTOM_HEADERS", Some("")),
|
||||
]);
|
||||
|
||||
let provider = create(
|
||||
"openai",
|
||||
ModelConfig::new_or_fail("gpt-4o-mini").with_canonical_limits("openai"),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(provider.as_lead_worker().is_none());
|
||||
assert_eq!(provider.get_model_config().model_name, "gpt-4o-mini");
|
||||
}
|
||||
|
||||
#[test_case::test_case(None, None, 16_000 ; "no overrides uses default")]
|
||||
#[test_case::test_case(Some("32000"), None, 32_000 ; "worker limit overrides default")]
|
||||
#[test_case::test_case(Some("32000"), Some("64000"), 32_000 ; "worker limit takes priority over global")]
|
||||
fn test_worker_model_context_limit(
|
||||
worker_limit: Option<&str>,
|
||||
global_limit: Option<&str>,
|
||||
expected_limit: usize,
|
||||
) {
|
||||
let _guard = env_lock::lock_env([
|
||||
("GOOSE_WORKER_CONTEXT_LIMIT", worker_limit),
|
||||
("GOOSE_CONTEXT_LIMIT", global_limit),
|
||||
]);
|
||||
|
||||
let default_model = ModelConfig::new_or_fail("gpt-3.5-turbo")
|
||||
.with_canonical_limits("openai")
|
||||
.with_context_limit(Some(16_000));
|
||||
|
||||
let result = create_worker_model_config(&default_model, "openai").unwrap();
|
||||
assert_eq!(result.context_limit, Some(expected_limit));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tanzu_declarative_provider_registry_wiring() {
|
||||
let providers_list = providers().await;
|
||||
|
||||
@@ -1,745 +0,0 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use super::base::{
|
||||
collect_stream, stream_from_single_message, LeadWorkerProviderTrait, MessageStream, Provider,
|
||||
ProviderDef, ProviderMetadata, ProviderUsage,
|
||||
};
|
||||
use super::errors::ProviderError;
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use crate::model::ModelConfig;
|
||||
use futures::future::BoxFuture;
|
||||
use rmcp::model::Tool;
|
||||
use rmcp::model::{Content, RawContent};
|
||||
|
||||
const LEAD_WORKER_PROVIDER_NAME: &str = "lead_worker";
|
||||
|
||||
/// A provider that switches between a lead model and a worker model based on turn count
|
||||
/// and can fall back to lead model on consecutive failures
|
||||
pub struct LeadWorkerProvider {
|
||||
lead_provider: Arc<dyn Provider>,
|
||||
worker_provider: Arc<dyn Provider>,
|
||||
lead_turns: usize,
|
||||
turn_count: Arc<Mutex<usize>>,
|
||||
failure_count: Arc<Mutex<usize>>,
|
||||
max_failures_before_fallback: usize,
|
||||
fallback_turns: usize,
|
||||
in_fallback_mode: Arc<Mutex<bool>>,
|
||||
fallback_remaining: Arc<Mutex<usize>>,
|
||||
}
|
||||
|
||||
impl LeadWorkerProvider {
|
||||
/// Create a new LeadWorkerProvider
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `lead_provider` - The provider to use for the initial turns
|
||||
/// * `worker_provider` - The provider to use after lead_turns
|
||||
/// * `lead_turns` - Number of turns to use the lead provider (default: 3)
|
||||
pub fn new(
|
||||
lead_provider: Arc<dyn Provider>,
|
||||
worker_provider: Arc<dyn Provider>,
|
||||
lead_turns: Option<usize>,
|
||||
) -> Self {
|
||||
Self {
|
||||
lead_provider,
|
||||
worker_provider,
|
||||
lead_turns: lead_turns.unwrap_or(3),
|
||||
turn_count: Arc::new(Mutex::new(0)),
|
||||
failure_count: Arc::new(Mutex::new(0)),
|
||||
max_failures_before_fallback: 2, // Fallback after 2 consecutive failures
|
||||
fallback_turns: 2, // Use lead model for 2 turns when in fallback mode
|
||||
in_fallback_mode: Arc::new(Mutex::new(false)),
|
||||
fallback_remaining: Arc::new(Mutex::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new LeadWorkerProvider with custom settings
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `lead_provider` - The provider to use for the initial turns
|
||||
/// * `worker_provider` - The provider to use after lead_turns
|
||||
/// * `lead_turns` - Number of turns to use the lead provider
|
||||
/// * `failure_threshold` - Number of consecutive failures before fallback
|
||||
/// * `fallback_turns` - Number of turns to use lead model in fallback mode
|
||||
pub fn new_with_settings(
|
||||
lead_provider: Arc<dyn Provider>,
|
||||
worker_provider: Arc<dyn Provider>,
|
||||
lead_turns: usize,
|
||||
failure_threshold: usize,
|
||||
fallback_turns: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
lead_provider,
|
||||
worker_provider,
|
||||
lead_turns,
|
||||
turn_count: Arc::new(Mutex::new(0)),
|
||||
failure_count: Arc::new(Mutex::new(0)),
|
||||
max_failures_before_fallback: failure_threshold,
|
||||
fallback_turns,
|
||||
in_fallback_mode: Arc::new(Mutex::new(false)),
|
||||
fallback_remaining: Arc::new(Mutex::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the turn counter and failure tracking (useful for new conversations)
|
||||
pub async fn reset_turn_count(&self) {
|
||||
let mut count = self.turn_count.lock().await;
|
||||
*count = 0;
|
||||
let mut failures = self.failure_count.lock().await;
|
||||
*failures = 0;
|
||||
let mut fallback = self.in_fallback_mode.lock().await;
|
||||
*fallback = false;
|
||||
let mut remaining = self.fallback_remaining.lock().await;
|
||||
*remaining = 0;
|
||||
}
|
||||
|
||||
/// Get the current turn count
|
||||
pub async fn get_turn_count(&self) -> usize {
|
||||
*self.turn_count.lock().await
|
||||
}
|
||||
|
||||
/// Get the current failure count
|
||||
pub async fn get_failure_count(&self) -> usize {
|
||||
*self.failure_count.lock().await
|
||||
}
|
||||
|
||||
/// Check if currently in fallback mode
|
||||
pub async fn is_in_fallback_mode(&self) -> bool {
|
||||
*self.in_fallback_mode.lock().await
|
||||
}
|
||||
|
||||
/// Get the currently active provider based on turn count and fallback state
|
||||
async fn get_active_provider(&self) -> Arc<dyn Provider> {
|
||||
let count = *self.turn_count.lock().await;
|
||||
let in_fallback = *self.in_fallback_mode.lock().await;
|
||||
|
||||
// Use lead provider if we're in initial turns OR in fallback mode
|
||||
if count < self.lead_turns || in_fallback {
|
||||
Arc::clone(&self.lead_provider)
|
||||
} else {
|
||||
Arc::clone(&self.worker_provider)
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle the result of a completion attempt and update failure tracking
|
||||
async fn handle_completion_result(
|
||||
&self,
|
||||
result: &Result<(Message, ProviderUsage), ProviderError>,
|
||||
) {
|
||||
match result {
|
||||
Ok((message, _usage)) => {
|
||||
// Check for task-level failures in the response
|
||||
let has_task_failure = self.detect_task_failures(message).await;
|
||||
|
||||
if has_task_failure {
|
||||
// Task failure detected - increment failure count
|
||||
let mut failures = self.failure_count.lock().await;
|
||||
*failures += 1;
|
||||
|
||||
let failure_count = *failures;
|
||||
let turn_count = *self.turn_count.lock().await;
|
||||
|
||||
tracing::warn!(
|
||||
"Task failure detected in response (failure count: {})",
|
||||
failure_count
|
||||
);
|
||||
|
||||
// Check if we should trigger fallback
|
||||
if turn_count >= self.lead_turns
|
||||
&& !*self.in_fallback_mode.lock().await
|
||||
&& failure_count >= self.max_failures_before_fallback
|
||||
{
|
||||
let mut in_fallback = self.in_fallback_mode.lock().await;
|
||||
let mut fallback_remaining = self.fallback_remaining.lock().await;
|
||||
|
||||
*in_fallback = true;
|
||||
*fallback_remaining = self.fallback_turns;
|
||||
*failures = 0; // Reset failure count when entering fallback
|
||||
|
||||
tracing::warn!(
|
||||
"🔄 SWITCHING TO LEAD MODEL: Entering fallback mode after {} consecutive task failures - using lead model for {} turns",
|
||||
self.max_failures_before_fallback,
|
||||
self.fallback_turns
|
||||
);
|
||||
}
|
||||
} else {
|
||||
// Success - reset failure count and handle fallback mode
|
||||
let mut failures = self.failure_count.lock().await;
|
||||
*failures = 0;
|
||||
|
||||
let mut in_fallback = self.in_fallback_mode.lock().await;
|
||||
let mut fallback_remaining = self.fallback_remaining.lock().await;
|
||||
|
||||
if *in_fallback {
|
||||
*fallback_remaining -= 1;
|
||||
if *fallback_remaining == 0 {
|
||||
*in_fallback = false;
|
||||
tracing::info!("✅ SWITCHING BACK TO WORKER MODEL: Exiting fallback mode - worker model resumed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Increment turn count on any completion (success or task failure)
|
||||
let mut count = self.turn_count.lock().await;
|
||||
*count += 1;
|
||||
}
|
||||
Err(_) => {
|
||||
// Technical failure - just log and let it bubble up
|
||||
// For technical failures (API/LLM issues), we don't want to second-guess
|
||||
// the model choice - just let the default model handle it
|
||||
tracing::warn!(
|
||||
"Technical failure detected - API/LLM issue, will use default model"
|
||||
);
|
||||
|
||||
// Don't increment turn count or failure tracking for technical failures
|
||||
// as these are temporary infrastructure issues, not model capability issues
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Detect task-level failures in the model's response
|
||||
async fn detect_task_failures(&self, message: &Message) -> bool {
|
||||
let mut failure_indicators = 0;
|
||||
|
||||
for content in &message.content {
|
||||
match content {
|
||||
MessageContent::ToolRequest(tool_request) => {
|
||||
// Check if tool request itself failed (malformed, etc.)
|
||||
if tool_request.tool_call.is_err() {
|
||||
failure_indicators += 1;
|
||||
tracing::debug!(
|
||||
"Failed tool request detected: {:?}",
|
||||
tool_request.tool_call
|
||||
);
|
||||
}
|
||||
}
|
||||
MessageContent::ToolResponse(tool_response) => {
|
||||
// Check if tool execution failed
|
||||
if let Err(tool_error) = &tool_response.tool_result {
|
||||
failure_indicators += 1;
|
||||
tracing::debug!("Tool execution failure detected: {:?}", tool_error);
|
||||
} else if let Ok(result) = &tool_response.tool_result {
|
||||
// Check tool output for error indicators
|
||||
if self.contains_error_indicators(&result.content) {
|
||||
failure_indicators += 1;
|
||||
tracing::debug!("Tool output contains error indicators");
|
||||
}
|
||||
}
|
||||
}
|
||||
MessageContent::Text(text_content) => {
|
||||
// Check for user correction patterns or error acknowledgments
|
||||
if self.contains_user_correction_patterns(&text_content.text) {
|
||||
failure_indicators += 1;
|
||||
tracing::debug!("User correction pattern detected in text");
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Consider it a failure if we have multiple failure indicators
|
||||
failure_indicators >= 1
|
||||
}
|
||||
|
||||
/// Check if tool output contains error indicators
|
||||
fn contains_error_indicators(&self, contents: &[Content]) -> bool {
|
||||
for content in contents {
|
||||
if let RawContent::Text(text_content) = content.deref() {
|
||||
let text_lower = text_content.text.to_lowercase();
|
||||
|
||||
// Common error patterns in tool outputs
|
||||
if text_lower.contains("error:")
|
||||
|| text_lower.contains("failed:")
|
||||
|| text_lower.contains("exception:")
|
||||
|| text_lower.contains("traceback")
|
||||
|| text_lower.contains("syntax error")
|
||||
|| text_lower.contains("permission denied")
|
||||
|| text_lower.contains("file not found")
|
||||
|| text_lower.contains("command not found")
|
||||
|| text_lower.contains("compilation failed")
|
||||
|| text_lower.contains("test failed")
|
||||
|| text_lower.contains("assertion failed")
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Check for user correction patterns in text
|
||||
fn contains_user_correction_patterns(&self, text: &str) -> bool {
|
||||
let text_lower = text.to_lowercase();
|
||||
|
||||
// Patterns indicating user is correcting or expressing dissatisfaction
|
||||
text_lower.contains("that's wrong")
|
||||
|| text_lower.contains("that's not right")
|
||||
|| text_lower.contains("that doesn't work")
|
||||
|| text_lower.contains("try again")
|
||||
|| text_lower.contains("let me correct")
|
||||
|| text_lower.contains("actually, ")
|
||||
|| text_lower.contains("no, that's")
|
||||
|| text_lower.contains("that's incorrect")
|
||||
|| text_lower.contains("fix this")
|
||||
|| text_lower.contains("this is broken")
|
||||
|| text_lower.contains("this doesn't")
|
||||
|| text_lower.starts_with("no,")
|
||||
|| text_lower.starts_with("wrong")
|
||||
|| text_lower.starts_with("incorrect")
|
||||
}
|
||||
}
|
||||
|
||||
impl LeadWorkerProviderTrait for LeadWorkerProvider {
|
||||
/// Get information about the lead and worker models for logging
|
||||
fn get_model_info(&self) -> (String, String) {
|
||||
let lead_model = self.lead_provider.get_model_config().model_name;
|
||||
let worker_model = self.worker_provider.get_model_config().model_name;
|
||||
(lead_model, worker_model)
|
||||
}
|
||||
|
||||
/// Get the currently active model name
|
||||
fn get_active_model(&self) -> String {
|
||||
// Read from the global store which was set during complete()
|
||||
use super::base::get_current_model;
|
||||
get_current_model().unwrap_or_else(|| {
|
||||
// Fallback to lead model if no current model is set
|
||||
self.lead_provider.get_model_config().model_name
|
||||
})
|
||||
}
|
||||
|
||||
/// Get (lead_turns, failure_threshold, fallback_turns)
|
||||
fn get_settings(&self) -> (usize, usize, usize) {
|
||||
(
|
||||
self.lead_turns,
|
||||
self.max_failures_before_fallback,
|
||||
self.fallback_turns,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderDef for LeadWorkerProvider {
|
||||
type Provider = Self;
|
||||
|
||||
fn metadata() -> ProviderMetadata {
|
||||
// This is a wrapper provider, so we return minimal metadata
|
||||
ProviderMetadata::new(
|
||||
LEAD_WORKER_PROVIDER_NAME,
|
||||
"Lead/Worker Provider",
|
||||
"A provider that switches between lead and worker models based on turn count",
|
||||
"", // No default model as this is determined by the wrapped providers
|
||||
vec![], // No known models as this depends on wrapped providers
|
||||
"", // No doc link
|
||||
vec![], // No config keys as configuration is done through wrapped providers
|
||||
)
|
||||
}
|
||||
|
||||
fn from_env(
|
||||
_model: ModelConfig,
|
||||
_extensions: Vec<crate::config::ExtensionConfig>,
|
||||
) -> BoxFuture<'static, Result<Self::Provider>> {
|
||||
Box::pin(async { Err(anyhow!("LeadWorkerProvider must be constructed explicitly")) })
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for LeadWorkerProvider {
|
||||
fn get_name(&self) -> &str {
|
||||
// Return the lead provider's name as the default
|
||||
self.lead_provider.get_name()
|
||||
}
|
||||
|
||||
fn get_model_config(&self) -> ModelConfig {
|
||||
// Return the lead provider's model config as the default
|
||||
// In practice, this might need to be more sophisticated
|
||||
self.lead_provider.get_model_config()
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
_model_config: &ModelConfig,
|
||||
session_id: &str,
|
||||
system: &str,
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
// Get the active provider
|
||||
let provider = self.get_active_provider().await;
|
||||
|
||||
// Log which provider is being used
|
||||
let turn_count = *self.turn_count.lock().await;
|
||||
let in_fallback = *self.in_fallback_mode.lock().await;
|
||||
let fallback_remaining = *self.fallback_remaining.lock().await;
|
||||
|
||||
let provider_type = if turn_count < self.lead_turns {
|
||||
"lead (initial)"
|
||||
} else if in_fallback {
|
||||
"lead (fallback)"
|
||||
} else {
|
||||
"worker"
|
||||
};
|
||||
|
||||
// Get the active model name and update the global store
|
||||
let active_model_name = if turn_count < self.lead_turns || in_fallback {
|
||||
self.lead_provider.get_model_config().model_name.clone()
|
||||
} else {
|
||||
self.worker_provider.get_model_config().model_name.clone()
|
||||
};
|
||||
|
||||
// Update the global current model store
|
||||
super::base::set_current_model(&active_model_name);
|
||||
|
||||
if in_fallback {
|
||||
tracing::info!(
|
||||
"🔄 Using {} provider for turn {} (FALLBACK MODE: {} turns remaining) - Model: {}",
|
||||
provider_type,
|
||||
turn_count + 1,
|
||||
fallback_remaining,
|
||||
active_model_name
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
"Using {} provider for turn {} (lead_turns: {}) - Model: {}",
|
||||
provider_type,
|
||||
turn_count + 1,
|
||||
self.lead_turns,
|
||||
active_model_name
|
||||
);
|
||||
}
|
||||
|
||||
// Make the completion request
|
||||
let model_config = provider.get_model_config();
|
||||
let stream_result = provider
|
||||
.stream(&model_config, session_id, system, messages, tools)
|
||||
.await;
|
||||
let result = match stream_result {
|
||||
Ok(stream) => collect_stream(stream).await,
|
||||
Err(e) => Err(e),
|
||||
};
|
||||
|
||||
// For technical failures, try with default model (lead provider) instead
|
||||
let final_result = match &result {
|
||||
Err(_) => {
|
||||
tracing::warn!("Technical failure with {} provider, retrying with default model (lead provider)", provider_type);
|
||||
|
||||
// Try with lead provider as the default/fallback for technical failures
|
||||
let model_config = self.lead_provider.get_model_config();
|
||||
let default_stream_result = self
|
||||
.lead_provider
|
||||
.stream(&model_config, session_id, system, messages, tools)
|
||||
.await;
|
||||
let default_result = match default_stream_result {
|
||||
Ok(stream) => collect_stream(stream).await,
|
||||
Err(e) => Err(e),
|
||||
};
|
||||
|
||||
match &default_result {
|
||||
Ok(_) => {
|
||||
tracing::info!(
|
||||
"✅ Default model (lead provider) succeeded after technical failure"
|
||||
);
|
||||
default_result
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::error!("❌ Default model (lead provider) also failed - returning original error");
|
||||
result // Return the original error
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(_) => result, // Success with original provider
|
||||
};
|
||||
|
||||
// Handle the result and update tracking (only for successful completions)
|
||||
self.handle_completion_result(&final_result).await;
|
||||
|
||||
match final_result {
|
||||
Ok((message, usage)) => Ok(stream_from_single_message(message, usage)),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
|
||||
// Combine models from both providers
|
||||
let mut all_models = self.lead_provider.fetch_supported_models().await?;
|
||||
let worker_models = self.worker_provider.fetch_supported_models().await?;
|
||||
all_models.extend(worker_models);
|
||||
all_models.sort();
|
||||
all_models.dedup();
|
||||
Ok(all_models)
|
||||
}
|
||||
|
||||
fn supports_embeddings(&self) -> bool {
|
||||
// Support embeddings if either provider supports them
|
||||
self.lead_provider.supports_embeddings() || self.worker_provider.supports_embeddings()
|
||||
}
|
||||
|
||||
async fn create_embeddings(
|
||||
&self,
|
||||
session_id: &str,
|
||||
texts: Vec<String>,
|
||||
) -> Result<Vec<Vec<f32>>, ProviderError> {
|
||||
// Use the lead provider for embeddings if it supports them, otherwise use worker
|
||||
if self.lead_provider.supports_embeddings() {
|
||||
self.lead_provider
|
||||
.create_embeddings(session_id, texts)
|
||||
.await
|
||||
} else if self.worker_provider.supports_embeddings() {
|
||||
self.worker_provider
|
||||
.create_embeddings(session_id, texts)
|
||||
.await
|
||||
} else {
|
||||
Err(ProviderError::ExecutionError(
|
||||
"Neither lead nor worker provider supports embeddings".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this provider is a LeadWorkerProvider
|
||||
fn as_lead_worker(&self) -> Option<&dyn LeadWorkerProviderTrait> {
|
||||
Some(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use crate::providers::base::{ProviderUsage, Usage};
|
||||
use chrono::Utc;
|
||||
use rmcp::model::{AnnotateAble, RawTextContent, Role};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockProvider {
|
||||
name: String,
|
||||
model_config: ModelConfig,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for MockProvider {
|
||||
fn get_name(&self) -> &str {
|
||||
"mock-lead"
|
||||
}
|
||||
|
||||
fn get_model_config(&self) -> ModelConfig {
|
||||
self.model_config.clone()
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
_model_config: &ModelConfig,
|
||||
_session_id: &str,
|
||||
_system: &str,
|
||||
_messages: &[Message],
|
||||
_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
Utc::now().timestamp(),
|
||||
vec![MessageContent::Text(
|
||||
RawTextContent {
|
||||
text: format!("Response from {}", self.name),
|
||||
meta: None,
|
||||
}
|
||||
.no_annotation(),
|
||||
)],
|
||||
);
|
||||
let usage = ProviderUsage::new(self.name.clone(), Usage::default());
|
||||
Ok(stream_from_single_message(message, usage))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_lead_worker_switching() {
|
||||
let lead_provider = Arc::new(MockProvider {
|
||||
name: "lead".to_string(),
|
||||
model_config: ModelConfig::new_or_fail("lead-model"),
|
||||
});
|
||||
|
||||
let worker_provider = Arc::new(MockProvider {
|
||||
name: "worker".to_string(),
|
||||
model_config: ModelConfig::new_or_fail("worker-model"),
|
||||
});
|
||||
|
||||
let provider = LeadWorkerProvider::new(lead_provider, worker_provider, Some(3));
|
||||
let model_config = provider.get_model_config();
|
||||
|
||||
// First three turns should use lead provider
|
||||
for i in 0..3 {
|
||||
let (_message, usage) = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(usage.model, "lead");
|
||||
assert_eq!(provider.get_turn_count().await, i + 1);
|
||||
assert!(!provider.is_in_fallback_mode().await);
|
||||
}
|
||||
|
||||
// Subsequent turns should use worker provider
|
||||
for i in 3..6 {
|
||||
let (_message, usage) = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(usage.model, "worker");
|
||||
assert_eq!(provider.get_turn_count().await, i + 1);
|
||||
assert!(!provider.is_in_fallback_mode().await);
|
||||
}
|
||||
|
||||
// Reset and verify it goes back to lead
|
||||
provider.reset_turn_count().await;
|
||||
assert_eq!(provider.get_turn_count().await, 0);
|
||||
assert_eq!(provider.get_failure_count().await, 0);
|
||||
assert!(!provider.is_in_fallback_mode().await);
|
||||
|
||||
let (_message, usage) = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(usage.model, "lead");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_technical_failure_retry() {
|
||||
let lead_provider = Arc::new(MockFailureProvider {
|
||||
name: "lead".to_string(),
|
||||
model_config: ModelConfig::new_or_fail("lead-model"),
|
||||
should_fail: false, // Lead provider works
|
||||
});
|
||||
|
||||
let worker_provider = Arc::new(MockFailureProvider {
|
||||
name: "worker".to_string(),
|
||||
model_config: ModelConfig::new_or_fail("worker-model"),
|
||||
should_fail: true, // Worker will fail
|
||||
});
|
||||
|
||||
let provider = LeadWorkerProvider::new(lead_provider, worker_provider, Some(2));
|
||||
let model_config = provider.get_model_config();
|
||||
|
||||
// First two turns use lead (should succeed)
|
||||
for _i in 0..2 {
|
||||
let result = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().1.model, "lead");
|
||||
assert!(!provider.is_in_fallback_mode().await);
|
||||
}
|
||||
|
||||
// Next turn uses worker (will fail, but should retry with lead and succeed)
|
||||
let model_config = provider.get_model_config();
|
||||
let result = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await;
|
||||
assert!(result.is_ok()); // Should succeed because lead provider is used as fallback
|
||||
assert_eq!(result.unwrap().1.model, "lead"); // Should be lead provider
|
||||
assert_eq!(provider.get_failure_count().await, 0); // No failure tracking for technical failures
|
||||
assert!(!provider.is_in_fallback_mode().await); // Not in fallback mode
|
||||
|
||||
// Another turn - should still try worker first, then retry with lead
|
||||
let model_config = provider.get_model_config();
|
||||
let result = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await;
|
||||
assert!(result.is_ok()); // Should succeed because lead provider is used as fallback
|
||||
assert_eq!(result.unwrap().1.model, "lead"); // Should be lead provider
|
||||
assert_eq!(provider.get_failure_count().await, 0); // Still no failure tracking
|
||||
assert!(!provider.is_in_fallback_mode().await); // Still not in fallback mode
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_fallback_on_task_failures() {
|
||||
// Test that task failures (not technical failures) still trigger fallback mode
|
||||
// This would need a different mock that simulates task failures in successful responses
|
||||
// For now, we'll test the fallback mode functionality directly
|
||||
let lead_provider = Arc::new(MockFailureProvider {
|
||||
name: "lead".to_string(),
|
||||
model_config: ModelConfig::new_or_fail("lead-model"),
|
||||
should_fail: false,
|
||||
});
|
||||
|
||||
let worker_provider = Arc::new(MockFailureProvider {
|
||||
name: "worker".to_string(),
|
||||
model_config: ModelConfig::new_or_fail("worker-model"),
|
||||
should_fail: false,
|
||||
});
|
||||
|
||||
let provider = LeadWorkerProvider::new(lead_provider, worker_provider, Some(2));
|
||||
|
||||
// Simulate being in fallback mode
|
||||
{
|
||||
let mut in_fallback = provider.in_fallback_mode.lock().await;
|
||||
*in_fallback = true;
|
||||
let mut fallback_remaining = provider.fallback_remaining.lock().await;
|
||||
*fallback_remaining = 2;
|
||||
let mut turn_count = provider.turn_count.lock().await;
|
||||
*turn_count = 4; // Past initial lead turns
|
||||
}
|
||||
|
||||
// Should use lead provider in fallback mode
|
||||
let model_config = provider.get_model_config();
|
||||
let result = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().1.model, "lead");
|
||||
assert!(provider.is_in_fallback_mode().await);
|
||||
|
||||
// One more fallback turn
|
||||
let model_config = provider.get_model_config();
|
||||
let result = provider
|
||||
.complete(&model_config, "test-session-id", "system", &[], &[])
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap().1.model, "lead");
|
||||
assert!(!provider.is_in_fallback_mode().await); // Should exit fallback mode
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockFailureProvider {
|
||||
name: String,
|
||||
model_config: ModelConfig,
|
||||
should_fail: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for MockFailureProvider {
|
||||
fn get_name(&self) -> &str {
|
||||
"mock-lead"
|
||||
}
|
||||
|
||||
fn get_model_config(&self) -> ModelConfig {
|
||||
self.model_config.clone()
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
_model_config: &ModelConfig,
|
||||
_session_id: &str,
|
||||
_system: &str,
|
||||
_messages: &[Message],
|
||||
_tools: &[Tool],
|
||||
) -> Result<MessageStream, ProviderError> {
|
||||
if self.should_fail {
|
||||
Err(ProviderError::ExecutionError(
|
||||
"Simulated failure".to_string(),
|
||||
))
|
||||
} else {
|
||||
let message = Message::new(
|
||||
Role::Assistant,
|
||||
Utc::now().timestamp(),
|
||||
vec![MessageContent::Text(
|
||||
RawTextContent {
|
||||
text: format!("Response from {}", self.name),
|
||||
meta: None,
|
||||
}
|
||||
.no_annotation(),
|
||||
)],
|
||||
);
|
||||
let usage = ProviderUsage::new(self.name.clone(), Usage::default());
|
||||
Ok(stream_from_single_message(message, usage))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -25,7 +25,6 @@ pub mod gemini_cli;
|
||||
pub mod githubcopilot;
|
||||
pub mod google;
|
||||
mod init;
|
||||
pub mod lead_worker;
|
||||
pub mod litellm;
|
||||
pub mod local_inference;
|
||||
pub mod nanogpt;
|
||||
|
||||
@@ -465,7 +465,6 @@ mod tests {
|
||||
responses.push(response);
|
||||
}
|
||||
Ok(AgentEvent::McpNotification(_)) => {}
|
||||
Ok(AgentEvent::ModelChange { .. }) => {}
|
||||
Ok(AgentEvent::HistoryReplaced(_updated_conversation)) => {
|
||||
// We should update the conversation here, but we're not reading it
|
||||
}
|
||||
|
||||
@@ -6113,28 +6113,6 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"required": [
|
||||
"model",
|
||||
"mode",
|
||||
"type"
|
||||
],
|
||||
"properties": {
|
||||
"mode": {
|
||||
"type": "string"
|
||||
},
|
||||
"model": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"ModelChange"
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "object",
|
||||
"required": [
|
||||
|
||||
@@ -687,10 +687,6 @@ export type MessageEvent = {
|
||||
reason: string;
|
||||
token_state: TokenState;
|
||||
type: 'Finish';
|
||||
} | {
|
||||
mode: string;
|
||||
model: string;
|
||||
type: 'ModelChange';
|
||||
} | {
|
||||
message: {
|
||||
[key: string]: unknown;
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
import { AppEvents } from '../constants/events';
|
||||
import React, {
|
||||
createContext,
|
||||
useCallback,
|
||||
useContext,
|
||||
useEffect,
|
||||
useMemo,
|
||||
useRef,
|
||||
@@ -42,8 +40,7 @@ import { useAutoSubmit } from '../hooks/useAutoSubmit';
|
||||
import { Goose } from './icons';
|
||||
import EnvironmentBadge from './GooseSidebar/EnvironmentBadge';
|
||||
|
||||
const CurrentModelContext = createContext<{ model: string; mode: string } | null>(null);
|
||||
export const useCurrentModelInfo = () => useContext(CurrentModelContext);
|
||||
|
||||
|
||||
interface BaseChatProps {
|
||||
setChat: (chat: ChatType) => void;
|
||||
|
||||
@@ -2,7 +2,6 @@ import { Sliders, Bot, Settings } from 'lucide-react';
|
||||
import React, { useEffect, useState } from 'react';
|
||||
import { useModelAndProvider } from '../../../ModelAndProviderContext';
|
||||
import { SwitchModelModal } from '../subcomponents/SwitchModelModal';
|
||||
import { LeadWorkerSettings } from '../subcomponents/LeadWorkerSettings';
|
||||
import { View } from '../../../../utils/navigationUtils';
|
||||
import {
|
||||
DropdownMenu,
|
||||
@@ -10,7 +9,6 @@ import {
|
||||
DropdownMenuItem,
|
||||
DropdownMenuTrigger,
|
||||
} from '../../../ui/dropdown-menu';
|
||||
import { useCurrentModelInfo } from '../../../BaseChat';
|
||||
import { useConfig } from '../../../ConfigContext';
|
||||
import { getProviderMetadata } from '../modelInterface';
|
||||
import { getModelDisplayName } from '../predefinedModelsUtils';
|
||||
@@ -46,81 +44,17 @@ export default function ModelsBottomBar({
|
||||
const currentModel = sessionModel ?? configModel;
|
||||
const currentProvider = sessionProvider ?? configProvider;
|
||||
|
||||
const currentModelInfo = useCurrentModelInfo();
|
||||
const { read, getProviders } = useConfig();
|
||||
const { getProviders } = useConfig();
|
||||
const [displayProvider, setDisplayProvider] = useState<string | null>(null);
|
||||
const [displayModelName, setDisplayModelName] = useState<string>('Select Model');
|
||||
const [isAddModelModalOpen, setIsAddModelModalOpen] = useState(false);
|
||||
const [isLeadWorkerModalOpen, setIsLeadWorkerModalOpen] = useState(false);
|
||||
const [isLocalModelSettingsOpen, setIsLocalModelSettingsOpen] = useState(false);
|
||||
const [isLeadWorkerActive, setIsLeadWorkerActive] = useState(false);
|
||||
const [providerDefaultModel, setProviderDefaultModel] = useState<string | null>(null);
|
||||
|
||||
// Check if lead/worker mode is active
|
||||
useEffect(() => {
|
||||
const checkLeadWorker = async () => {
|
||||
try {
|
||||
const leadModel = await read('GOOSE_LEAD_MODEL', false);
|
||||
setIsLeadWorkerActive(!!leadModel);
|
||||
} catch (error) {
|
||||
console.error('Error checking lead model:', error);
|
||||
setIsLeadWorkerActive(false);
|
||||
}
|
||||
};
|
||||
checkLeadWorker();
|
||||
}, [read]);
|
||||
|
||||
// Refresh lead/worker status when modal closes
|
||||
const handleLeadWorkerModalClose = () => {
|
||||
setIsLeadWorkerModalOpen(false);
|
||||
const checkLeadWorker = async () => {
|
||||
try {
|
||||
const leadModel = await read('GOOSE_LEAD_MODEL', false);
|
||||
const currentModel = await read('GOOSE_MODEL', false);
|
||||
setIsLeadWorkerActive(!!leadModel);
|
||||
setLeadModelName((leadModel as string) || '');
|
||||
setCurrentActiveModel((currentModel as string) || '');
|
||||
} catch (error) {
|
||||
console.error('Error checking lead model after modal close:', error);
|
||||
setIsLeadWorkerActive(false);
|
||||
}
|
||||
};
|
||||
checkLeadWorker();
|
||||
};
|
||||
|
||||
const [leadModelName, setLeadModelName] = useState<string>('');
|
||||
const [currentActiveModel, setCurrentActiveModel] = useState<string>('');
|
||||
|
||||
// Get lead model name and current model for comparison
|
||||
useEffect(() => {
|
||||
const getModelInfo = async () => {
|
||||
try {
|
||||
const leadModel = await read('GOOSE_LEAD_MODEL', false);
|
||||
const currentModel = await read('GOOSE_MODEL', false);
|
||||
setLeadModelName((leadModel as string) || '');
|
||||
setCurrentActiveModel((currentModel as string) || '');
|
||||
} catch (error) {
|
||||
console.error('Error getting model info:', error);
|
||||
}
|
||||
};
|
||||
getModelInfo();
|
||||
}, [read]);
|
||||
|
||||
// Determine the mode based on which model is currently active
|
||||
const modelMode = isLeadWorkerActive
|
||||
? currentActiveModel === leadModelName
|
||||
? 'lead'
|
||||
: 'worker'
|
||||
: undefined;
|
||||
|
||||
// Determine which model to display - activeModel takes priority when lead/worker is active
|
||||
// Hide label while session data is still being fetched (avoids flashing
|
||||
// the config default before the session's actual model arrives).
|
||||
const isModelLoading = sessionId && !sessionLoaded;
|
||||
const displayModel =
|
||||
isLeadWorkerActive && currentModelInfo?.model
|
||||
? currentModelInfo.model
|
||||
: currentModel || providerDefaultModel || displayModelName;
|
||||
const displayModel = currentModel || providerDefaultModel || displayModelName;
|
||||
|
||||
useEffect(() => {
|
||||
if (!currentProvider) return;
|
||||
@@ -168,9 +102,6 @@ export default function ModelsBottomBar({
|
||||
<Bot className="mr-1 h-4 w-4 flex-shrink-0" />
|
||||
<span className={`truncate text-xs${isModelLoading ? ' opacity-0' : ''}`}>
|
||||
{displayModel}
|
||||
{isLeadWorkerActive && modelMode && (
|
||||
<span className="ml-1 text-[10px] opacity-60">({modelMode})</span>
|
||||
)}
|
||||
</span>
|
||||
</div>
|
||||
</DropdownMenuTrigger>
|
||||
@@ -184,10 +115,6 @@ export default function ModelsBottomBar({
|
||||
<span>Change Model</span>
|
||||
<Sliders className="ml-auto h-4 w-4 rotate-90" />
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem onClick={() => setIsLeadWorkerModalOpen(true)}>
|
||||
<span>Lead/Worker Settings</span>
|
||||
<Sliders className="ml-auto h-4 w-4" />
|
||||
</DropdownMenuItem>
|
||||
{currentProvider === 'local' && currentModel && (
|
||||
<DropdownMenuItem onClick={() => setIsLocalModelSettingsOpen(true)}>
|
||||
<span>Local Model Settings</span>
|
||||
@@ -208,10 +135,6 @@ export default function ModelsBottomBar({
|
||||
/>
|
||||
) : null}
|
||||
|
||||
{isLeadWorkerModalOpen ? (
|
||||
<LeadWorkerSettings isOpen={isLeadWorkerModalOpen} onClose={handleLeadWorkerModalClose} />
|
||||
) : null}
|
||||
|
||||
{isLocalModelSettingsOpen && currentModel && (
|
||||
<div className="fixed inset-0 z-50 flex items-center justify-center bg-black/50">
|
||||
<div className="bg-background-default rounded-lg shadow-lg w-[480px] max-h-[80vh] flex flex-col">
|
||||
|
||||
@@ -1,147 +0,0 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { render, screen, waitFor, fireEvent } from '@testing-library/react';
|
||||
import { LeadWorkerSettings } from './LeadWorkerSettings';
|
||||
|
||||
// Mock predefined models utils to force provider-based options (no predefined list)
|
||||
vi.mock('../predefinedModelsUtils', () => ({
|
||||
shouldShowPredefinedModels: () => false,
|
||||
getPredefinedModelsFromEnv: () => [],
|
||||
}));
|
||||
|
||||
// Mocks for useConfig
|
||||
const mockRead = vi.fn();
|
||||
const mockUpsert = vi.fn();
|
||||
const mockRemove = vi.fn();
|
||||
const mockGetProviders = vi.fn();
|
||||
|
||||
vi.mock('../../../ConfigContext', () => ({
|
||||
useConfig: () => ({
|
||||
read: mockRead,
|
||||
upsert: mockUpsert,
|
||||
remove: mockRemove,
|
||||
getProviders: mockGetProviders,
|
||||
}),
|
||||
}));
|
||||
|
||||
describe('LeadWorkerSettings', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
const setupHappyPathMocks = () => {
|
||||
// reads
|
||||
mockRead.mockImplementation(async (key: string) => {
|
||||
switch (key) {
|
||||
case 'GOOSE_LEAD_MODEL':
|
||||
return 'my-custom-lead';
|
||||
case 'GOOSE_LEAD_PROVIDER':
|
||||
return 'anthropic';
|
||||
case 'GOOSE_LEAD_TURNS':
|
||||
return 3;
|
||||
case 'GOOSE_LEAD_FAILURE_THRESHOLD':
|
||||
return 2;
|
||||
case 'GOOSE_LEAD_FALLBACK_TURNS':
|
||||
return 2;
|
||||
case 'GOOSE_MODEL':
|
||||
return 'my-custom-worker';
|
||||
case 'GOOSE_PROVIDER':
|
||||
return 'openai';
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
});
|
||||
|
||||
// providers (options do NOT include the custom models above)
|
||||
mockGetProviders.mockResolvedValue([
|
||||
{
|
||||
is_configured: true,
|
||||
name: 'openai',
|
||||
metadata: {
|
||||
display_name: 'OpenAI',
|
||||
known_models: [{ name: 'gpt-4o' }, { name: 'gpt-4o-mini' }],
|
||||
},
|
||||
},
|
||||
{
|
||||
is_configured: true,
|
||||
name: 'anthropic',
|
||||
metadata: {
|
||||
display_name: 'Anthropic',
|
||||
known_models: [{ name: 'claude-3-5-sonnet' }],
|
||||
},
|
||||
},
|
||||
]);
|
||||
|
||||
// writers
|
||||
mockUpsert.mockResolvedValue(undefined);
|
||||
mockRemove.mockResolvedValue(undefined);
|
||||
};
|
||||
|
||||
it('shows custom inputs for lead/worker when current models are unknown and saves them', async () => {
|
||||
setupHappyPathMocks();
|
||||
|
||||
const onClose = vi.fn();
|
||||
render(<LeadWorkerSettings isOpen={true} onClose={onClose} />);
|
||||
|
||||
// Wait for modal content (not loading)
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText('Lead/Worker Mode')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Labels should be present with back-to-list controls
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText('Lead Model')).toBeInTheDocument();
|
||||
expect(screen.getByText('Worker Model')).toBeInTheDocument();
|
||||
// Back to model list appears for each section when in custom mode
|
||||
const backLinks = screen.getAllByText('Back to model list');
|
||||
expect(backLinks.length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
|
||||
const inputs = screen.getAllByPlaceholderText('Type model name here') as HTMLInputElement[];
|
||||
expect(inputs.length).toBe(2);
|
||||
const [leadInput, workerInput] = inputs;
|
||||
expect(leadInput.value).toBe('my-custom-lead');
|
||||
expect(workerInput.value).toBe('my-custom-worker');
|
||||
|
||||
// Save settings
|
||||
const saveBtn = screen.getByRole('button', { name: 'Save Settings' });
|
||||
expect(saveBtn).toBeEnabled();
|
||||
fireEvent.click(saveBtn);
|
||||
|
||||
// Assert upserts for models (providers are optional but present in this setup)
|
||||
await waitFor(() => {
|
||||
expect(mockUpsert).toHaveBeenCalledWith('GOOSE_LEAD_MODEL', 'my-custom-lead', false);
|
||||
expect(mockUpsert).toHaveBeenCalledWith('GOOSE_MODEL', 'my-custom-worker', false);
|
||||
expect(mockUpsert).toHaveBeenCalledWith('GOOSE_LEAD_PROVIDER', 'anthropic', false);
|
||||
expect(mockUpsert).toHaveBeenCalledWith('GOOSE_PROVIDER', 'openai', false);
|
||||
});
|
||||
});
|
||||
|
||||
it('disables lead/worker and removes config when toggled off', async () => {
|
||||
setupHappyPathMocks();
|
||||
|
||||
const onClose = vi.fn();
|
||||
render(<LeadWorkerSettings isOpen={true} onClose={onClose} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText('Lead/Worker Mode')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
// Toggle off
|
||||
const checkbox = screen.getByLabelText('Enable lead/worker mode') as HTMLInputElement;
|
||||
expect(checkbox.checked).toBe(true);
|
||||
fireEvent.click(checkbox);
|
||||
expect(checkbox.checked).toBe(false);
|
||||
|
||||
const saveBtn = screen.getByRole('button', { name: 'Save Settings' });
|
||||
expect(saveBtn).toBeEnabled();
|
||||
fireEvent.click(saveBtn);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_MODEL', false);
|
||||
expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_PROVIDER', false);
|
||||
expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_TURNS', false);
|
||||
expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_FAILURE_THRESHOLD', false);
|
||||
expect(mockRemove).toHaveBeenCalledWith('GOOSE_LEAD_FALLBACK_TURNS', false);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,419 +0,0 @@
|
||||
import { useState, useEffect } from 'react';
|
||||
import { useConfig } from '../../../ConfigContext';
|
||||
import { Button } from '../../../ui/button';
|
||||
import { Select } from '../../../ui/Select';
|
||||
import { Input } from '../../../ui/input';
|
||||
import { getPredefinedModelsFromEnv, shouldShowPredefinedModels } from '../predefinedModelsUtils';
|
||||
import { Dialog, DialogContent, DialogHeader, DialogTitle } from '../../../ui/dialog';
|
||||
import { fetchModelsForProviders } from '../modelInterface';
|
||||
|
||||
interface LeadWorkerSettingsProps {
|
||||
isOpen: boolean;
|
||||
onClose: () => void;
|
||||
}
|
||||
|
||||
export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps) {
|
||||
const { read, upsert, getProviders, remove } = useConfig();
|
||||
const [leadModel, setLeadModel] = useState<string>('');
|
||||
const [workerModel, setWorkerModel] = useState<string>('');
|
||||
const [leadProvider, setLeadProvider] = useState<string>('');
|
||||
const [workerProvider, setWorkerProvider] = useState<string>('');
|
||||
// Minimal custom model mode toggles
|
||||
const [isLeadCustomModel, setIsLeadCustomModel] = useState<boolean>(false);
|
||||
const [isWorkerCustomModel, setIsWorkerCustomModel] = useState<boolean>(false);
|
||||
const [leadTurns, setLeadTurns] = useState<number>(3);
|
||||
const [failureThreshold, setFailureThreshold] = useState<number>(2);
|
||||
const [fallbackTurns, setFallbackTurns] = useState<number>(2);
|
||||
const [isEnabled, setIsEnabled] = useState(false);
|
||||
const [modelOptions, setModelOptions] = useState<
|
||||
{ value: string; label: string; provider: string }[]
|
||||
>([]);
|
||||
const [isLoading, setIsLoading] = useState(true);
|
||||
|
||||
// Load current configuration
|
||||
useEffect(() => {
|
||||
if (!isOpen) return; // Only load when modal is open
|
||||
|
||||
const loadConfig = async () => {
|
||||
try {
|
||||
setIsLoading(true);
|
||||
const [
|
||||
leadModelConfig,
|
||||
leadProviderConfig,
|
||||
leadTurnsConfig,
|
||||
failureThresholdConfig,
|
||||
fallbackTurnsConfig,
|
||||
] = await Promise.all([
|
||||
read('GOOSE_LEAD_MODEL', false),
|
||||
read('GOOSE_LEAD_PROVIDER', false),
|
||||
read('GOOSE_LEAD_TURNS', false),
|
||||
read('GOOSE_LEAD_FAILURE_THRESHOLD', false),
|
||||
read('GOOSE_LEAD_FALLBACK_TURNS', false),
|
||||
]);
|
||||
|
||||
if (leadModelConfig) {
|
||||
setLeadModel(leadModelConfig as string);
|
||||
setIsEnabled(true);
|
||||
} else {
|
||||
setLeadModel('');
|
||||
setIsEnabled(false);
|
||||
}
|
||||
if (leadProviderConfig) setLeadProvider(leadProviderConfig as string);
|
||||
else setLeadProvider('');
|
||||
if (leadTurnsConfig) setLeadTurns(Number(leadTurnsConfig));
|
||||
else setLeadTurns(3);
|
||||
if (failureThresholdConfig) setFailureThreshold(Number(failureThresholdConfig));
|
||||
else setFailureThreshold(2);
|
||||
if (fallbackTurnsConfig) setFallbackTurns(Number(fallbackTurnsConfig));
|
||||
else setFallbackTurns(2);
|
||||
|
||||
// Set worker model from config
|
||||
const workerModelConfig = await read('GOOSE_MODEL', false);
|
||||
if (workerModelConfig) {
|
||||
setWorkerModel(workerModelConfig as string);
|
||||
} else {
|
||||
setWorkerModel('');
|
||||
}
|
||||
|
||||
const workerProviderConfig = await read('GOOSE_PROVIDER', false);
|
||||
if (workerProviderConfig) {
|
||||
setWorkerProvider(workerProviderConfig as string);
|
||||
} else {
|
||||
setWorkerProvider('');
|
||||
}
|
||||
|
||||
// Load available models
|
||||
const options: { value: string; label: string; provider: string }[] = [];
|
||||
|
||||
if (shouldShowPredefinedModels()) {
|
||||
// Use predefined models if available
|
||||
const predefinedModels = getPredefinedModelsFromEnv();
|
||||
predefinedModels.forEach((model) => {
|
||||
options.push({
|
||||
value: model.name, // Use name for switching
|
||||
label: model.alias || model.name, // Use alias for display, fall back to name
|
||||
provider: model.provider,
|
||||
});
|
||||
});
|
||||
} else {
|
||||
// Fallback to provider-based models
|
||||
const providers = await getProviders(false);
|
||||
const activeProviders = providers.filter((p) => p.is_configured);
|
||||
|
||||
const results = await fetchModelsForProviders(activeProviders);
|
||||
|
||||
results.forEach(({ provider: p, models, error }) => {
|
||||
if (error) {
|
||||
console.error(error);
|
||||
}
|
||||
|
||||
if (models && models.length > 0) {
|
||||
models.forEach((modelName) => {
|
||||
options.push({
|
||||
value: modelName,
|
||||
label: `${modelName} (${p.metadata.display_name})`,
|
||||
provider: p.name,
|
||||
});
|
||||
});
|
||||
}
|
||||
// Add custom model option for all non-Custom providers
|
||||
if (p.provider_type !== 'Custom') {
|
||||
options.push({
|
||||
value: `__custom__:${p.name}`,
|
||||
label: 'Enter a model not listed...',
|
||||
provider: p.name,
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
setModelOptions(options);
|
||||
} catch (error) {
|
||||
console.error('Error loading configuration:', error);
|
||||
} finally {
|
||||
setIsLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
loadConfig();
|
||||
}, [read, getProviders, isOpen]);
|
||||
|
||||
// If current models are not in the list (e.g., previously set to custom), switch to custom mode
|
||||
useEffect(() => {
|
||||
if (!isLoading) {
|
||||
if (leadModel && !modelOptions.find((opt) => opt.value === leadModel)) {
|
||||
setIsLeadCustomModel(true);
|
||||
}
|
||||
if (workerModel && !modelOptions.find((opt) => opt.value === workerModel)) {
|
||||
setIsWorkerCustomModel(true);
|
||||
}
|
||||
}
|
||||
}, [isLoading, modelOptions, leadModel, workerModel]);
|
||||
|
||||
const handleSave = async () => {
|
||||
try {
|
||||
if (isEnabled && leadModel && workerModel) {
|
||||
// Save lead/worker configuration
|
||||
await Promise.all([
|
||||
upsert('GOOSE_LEAD_MODEL', leadModel, false),
|
||||
leadProvider && upsert('GOOSE_LEAD_PROVIDER', leadProvider, false),
|
||||
upsert('GOOSE_MODEL', workerModel, false),
|
||||
workerProvider && upsert('GOOSE_PROVIDER', workerProvider, false),
|
||||
upsert('GOOSE_LEAD_TURNS', leadTurns, false),
|
||||
upsert('GOOSE_LEAD_FAILURE_THRESHOLD', failureThreshold, false),
|
||||
upsert('GOOSE_LEAD_FALLBACK_TURNS', fallbackTurns, false),
|
||||
]);
|
||||
} else {
|
||||
// Remove lead/worker configuration
|
||||
await Promise.all([
|
||||
remove('GOOSE_LEAD_MODEL', false),
|
||||
remove('GOOSE_LEAD_PROVIDER', false),
|
||||
remove('GOOSE_LEAD_TURNS', false),
|
||||
remove('GOOSE_LEAD_FAILURE_THRESHOLD', false),
|
||||
remove('GOOSE_LEAD_FALLBACK_TURNS', false),
|
||||
]);
|
||||
}
|
||||
onClose();
|
||||
} catch (error) {
|
||||
console.error('Error saving configuration:', error);
|
||||
}
|
||||
};
|
||||
|
||||
if (isLoading) {
|
||||
return (
|
||||
<Dialog open={isOpen} onOpenChange={(open) => !open && onClose()}>
|
||||
<DialogContent className="sm:max-w-[500px]">
|
||||
<DialogHeader>
|
||||
<DialogTitle>Lead/Worker Mode</DialogTitle>
|
||||
</DialogHeader>
|
||||
<div className="p-4">Loading...</div>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<Dialog open={isOpen} onOpenChange={(open) => !open && onClose()}>
|
||||
<DialogContent className="sm:max-w-[500px]">
|
||||
<DialogHeader>
|
||||
<DialogTitle>Lead/Worker Mode</DialogTitle>
|
||||
</DialogHeader>
|
||||
<div className="p-4 space-y-4">
|
||||
<div className="space-y-2">
|
||||
<p className="text-sm text-text-secondary">
|
||||
Configure a lead model for planning and a worker model for execution
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center space-x-2">
|
||||
<input
|
||||
type="checkbox"
|
||||
id="enable-lead-worker"
|
||||
checked={isEnabled}
|
||||
onChange={(e) => setIsEnabled(e.target.checked)}
|
||||
className="rounded border-border-primary"
|
||||
/>
|
||||
<label htmlFor="enable-lead-worker" className="text-sm text-text-primary">
|
||||
Enable lead/worker mode
|
||||
</label>
|
||||
</div>
|
||||
|
||||
<div className="space-y-4">
|
||||
<div className="space-y-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<label
|
||||
className={`text-sm ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Lead Model
|
||||
</label>
|
||||
{isLeadCustomModel && (
|
||||
<button
|
||||
onClick={() => setIsLeadCustomModel(false)}
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'} hover:underline`}
|
||||
type="button"
|
||||
>
|
||||
Back to model list
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
{!isLeadCustomModel ? (
|
||||
<Select
|
||||
options={modelOptions}
|
||||
value={
|
||||
leadModel ? modelOptions.find((opt) => opt.value === leadModel) || null : null
|
||||
}
|
||||
onChange={(newValue: unknown) => {
|
||||
const option = newValue as { value: string; provider: string } | null;
|
||||
if (option) {
|
||||
if (option.value.startsWith('__custom__')) {
|
||||
setIsLeadCustomModel(true);
|
||||
setLeadModel('');
|
||||
setLeadProvider(option.provider);
|
||||
return;
|
||||
}
|
||||
setLeadModel(option.value);
|
||||
setLeadProvider(option.provider);
|
||||
}
|
||||
}}
|
||||
placeholder="Select lead model..."
|
||||
isDisabled={!isEnabled}
|
||||
className={!isEnabled ? 'opacity-50' : ''}
|
||||
/>
|
||||
) : (
|
||||
<Input
|
||||
className="h-[38px] mb-2"
|
||||
placeholder="Type model name here"
|
||||
onChange={(event) => setLeadModel(event.target.value)}
|
||||
value={leadModel}
|
||||
disabled={!isEnabled}
|
||||
/>
|
||||
)}
|
||||
<p
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Strong model for initial planning and fallback recovery
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div className="space-y-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<label
|
||||
className={`text-sm ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Worker Model
|
||||
</label>
|
||||
{isWorkerCustomModel && (
|
||||
<button
|
||||
onClick={() => setIsWorkerCustomModel(false)}
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'} hover:underline`}
|
||||
type="button"
|
||||
>
|
||||
Back to model list
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
{!isWorkerCustomModel ? (
|
||||
<Select
|
||||
options={modelOptions}
|
||||
value={
|
||||
workerModel
|
||||
? modelOptions.find((opt) => opt.value === workerModel) || null
|
||||
: null
|
||||
}
|
||||
onChange={(newValue: unknown) => {
|
||||
const option = newValue as { value: string; provider: string } | null;
|
||||
if (option) {
|
||||
if (option.value.startsWith('__custom__')) {
|
||||
setIsWorkerCustomModel(true);
|
||||
setWorkerModel('');
|
||||
setWorkerProvider(option.provider);
|
||||
return;
|
||||
}
|
||||
setWorkerModel(option.value);
|
||||
setWorkerProvider(option.provider);
|
||||
}
|
||||
}}
|
||||
placeholder="Select worker model..."
|
||||
isDisabled={!isEnabled}
|
||||
className={!isEnabled ? 'opacity-50' : ''}
|
||||
/>
|
||||
) : (
|
||||
<Input
|
||||
className="h-[38px] mb-2"
|
||||
placeholder="Type model name here"
|
||||
onChange={(event) => setWorkerModel(event.target.value)}
|
||||
value={workerModel}
|
||||
disabled={!isEnabled}
|
||||
/>
|
||||
)}
|
||||
<p
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Fast model for routine execution tasks
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div
|
||||
className={`space-y-4 pt-4 border-t border-border-primary ${!isEnabled ? 'opacity-50' : ''}`}
|
||||
>
|
||||
<div className="space-y-2">
|
||||
<label
|
||||
className={`text-sm flex items-center gap-1 ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Initial Lead Turns
|
||||
</label>
|
||||
<Input
|
||||
type="number"
|
||||
min={1}
|
||||
max={10}
|
||||
value={leadTurns}
|
||||
onChange={(e) => setLeadTurns(Number(e.target.value))}
|
||||
className={`w-20 ${!isEnabled ? 'opacity-50 cursor-not-allowed' : ''}`}
|
||||
disabled={!isEnabled}
|
||||
/>
|
||||
<p
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Number of turns to use the lead model at the start
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div className="space-y-2">
|
||||
<label
|
||||
className={`text-sm flex items-center gap-1 ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Failure Threshold
|
||||
</label>
|
||||
<Input
|
||||
type="number"
|
||||
min={1}
|
||||
max={5}
|
||||
value={failureThreshold}
|
||||
onChange={(e) => setFailureThreshold(Number(e.target.value))}
|
||||
className={`w-20 ${!isEnabled ? 'opacity-50 cursor-not-allowed' : ''}`}
|
||||
disabled={!isEnabled}
|
||||
/>
|
||||
<p
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Consecutive failures before switching back to lead
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div className="space-y-2">
|
||||
<label
|
||||
className={`text-sm flex items-center gap-1 ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Fallback Turns
|
||||
</label>
|
||||
<Input
|
||||
type="number"
|
||||
min={1}
|
||||
max={5}
|
||||
value={fallbackTurns}
|
||||
onChange={(e) => setFallbackTurns(Number(e.target.value))}
|
||||
className={`w-20 ${!isEnabled ? 'opacity-50 cursor-not-allowed' : ''}`}
|
||||
disabled={!isEnabled}
|
||||
/>
|
||||
<p
|
||||
className={`text-xs ${!isEnabled ? 'text-text-secondary' : 'text-text-secondary'}`}
|
||||
>
|
||||
Turns to use lead model during fallback
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="flex justify-end space-x-2 pt-4 border-t border-border-primary">
|
||||
<Button variant="ghost" onClick={onClose}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button onClick={handleSave} disabled={isEnabled && (!leadModel || !workerModel)}>
|
||||
Save Settings
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -318,9 +318,6 @@ function createEventProcessor(
|
||||
onFinish();
|
||||
return true;
|
||||
}
|
||||
case 'ModelChange': {
|
||||
return false;
|
||||
}
|
||||
case 'UpdateConversation': {
|
||||
const conversation = (event as Record<string, unknown>).conversation as Message[];
|
||||
currentMessages = conversation;
|
||||
|
||||
@@ -4,8 +4,6 @@ export const configLabels: Record<string, string> = {
|
||||
GOOSE_MODEL: 'Model',
|
||||
GOOSE_TEMPERATURE: 'Temperature',
|
||||
GOOSE_MODE: 'Mode',
|
||||
GOOSE_LEAD_PROVIDER: 'Lead Provider',
|
||||
GOOSE_LEAD_MODEL: 'Lead Model',
|
||||
GOOSE_PLANNER_PROVIDER: 'Planner Provider',
|
||||
GOOSE_PLANNER_MODEL: 'Planner Model',
|
||||
GOOSE_TOOLSHIM: 'Tool Shim',
|
||||
|
||||
Reference in New Issue
Block a user