From fbe46d0b22b7039bc4ef6f093d12d6f5c450bfd3 Mon Sep 17 00:00:00 2001 From: Rabi Mishra Date: Tue, 20 Jan 2026 22:19:10 +0530 Subject: [PATCH] feat(gcp-vertex): add model list with org policy filtering (#6393) Signed-off-by: rabi --- crates/goose-cli/src/commands/configure.rs | 60 ++++- crates/goose/src/providers/base.rs | 12 + .../src/providers/formats/gcpvertexai.rs | 226 +++++------------- crates/goose/src/providers/gcpvertexai.rs | 146 ++++++++--- .../goose/src/providers/provider_registry.rs | 1 + crates/goose/tests/agent.rs | 1 + ui/desktop/openapi.json | 4 + ui/desktop/src/api/types.gen.ts | 4 + .../subcomponents/LeadWorkerSettings.tsx | 17 +- .../models/subcomponents/SwitchModelModal.tsx | 59 ++--- 10 files changed, 283 insertions(+), 247 deletions(-) diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index c870ed4b19..8247221c76 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -436,12 +436,12 @@ fn select_model_from_list( provider_meta: &goose::providers::base::ProviderMetadata, ) -> anyhow::Result { const MAX_MODELS: usize = 10; + const UNLISTED_MODEL_KEY: &str = "__unlisted__"; + // Smart model selection: // If we have more than MAX_MODELS models, show the recommended models with additional search option. // Otherwise, show all models without search. - if models.len() > MAX_MODELS { - // Get recommended models from provider metadata let recommended_models: Vec = provider_meta .known_models .iter() @@ -464,12 +464,22 @@ fn select_model_from_list( ), ); + if provider_meta.allows_unlisted_models { + model_items.push(( + UNLISTED_MODEL_KEY.to_string(), + "Enter a model not listed...".to_string(), + "", + )); + } + let selection = cliclack::select("Select a model:") .items(&model_items) .interact()?; if selection == "search_all" { Ok(interactive_model_search(models)?) + } else if selection == UNLISTED_MODEL_KEY { + prompt_unlisted_model(provider_meta) } else { Ok(selection) } @@ -477,19 +487,45 @@ fn select_model_from_list( Ok(interactive_model_search(models)?) } } else { - // just a few models, show all without search for better UX - Ok(cliclack::select("Select a model:") - .items( - &models - .iter() - .map(|m| (m, m.as_str(), "")) - .collect::>(), - ) - .interact()? - .to_string()) + let mut model_items: Vec<(String, String, &str)> = + models.iter().map(|m| (m.clone(), m.clone(), "")).collect(); + + if provider_meta.allows_unlisted_models { + model_items.push(( + UNLISTED_MODEL_KEY.to_string(), + "Enter a model not listed...".to_string(), + "", + )); + } + + let selection = cliclack::select("Select a model:") + .items(&model_items) + .interact()?; + + if selection == UNLISTED_MODEL_KEY { + prompt_unlisted_model(provider_meta) + } else { + Ok(selection) + } } } +fn prompt_unlisted_model( + provider_meta: &goose::providers::base::ProviderMetadata, +) -> anyhow::Result { + let model: String = cliclack::input("Enter the model name:") + .placeholder(&provider_meta.default_model) + .validate(|input: &String| { + if input.trim().is_empty() { + Err("Please enter a model name") + } else { + Ok(()) + } + }) + .interact()?; + Ok(model.trim().to_string()) +} + fn try_store_secret(config: &Config, key_name: &str, value: String) -> anyhow::Result { match config.set_secret(key_name, &value) { Ok(_) => Ok(true), diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index cec5734d94..86c69828c2 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -108,6 +108,9 @@ pub struct ProviderMetadata { pub model_doc_link: String, /// Required configuration keys pub config_keys: Vec, + /// Whether this provider allows entering model names not in the fetched list + #[serde(default)] + pub allows_unlisted_models: bool, } impl ProviderMetadata { @@ -138,6 +141,7 @@ impl ProviderMetadata { .collect(), model_doc_link: model_doc_link.to_string(), config_keys, + allows_unlisted_models: false, } } @@ -158,6 +162,7 @@ impl ProviderMetadata { known_models: models, model_doc_link: model_doc_link.to_string(), config_keys, + allows_unlisted_models: false, } } @@ -170,8 +175,15 @@ impl ProviderMetadata { known_models: vec![], model_doc_link: "".to_string(), config_keys: vec![], + allows_unlisted_models: false, } } + + /// Set allows_unlisted_models flag (builder pattern) + pub fn with_unlisted_models(mut self) -> Self { + self.allows_unlisted_models = true; + self + } } /// Configuration key metadata for provider setup diff --git a/crates/goose/src/providers/formats/gcpvertexai.rs b/crates/goose/src/providers/formats/gcpvertexai.rs index ae1f306f57..16edea891d 100644 --- a/crates/goose/src/providers/formats/gcpvertexai.rs +++ b/crates/goose/src/providers/formats/gcpvertexai.rs @@ -70,77 +70,49 @@ pub enum ModelError { UnsupportedLocation(String), } +/// Default model for GCP Vertex AI. +pub const DEFAULT_MODEL: &str = "gemini-2.5-flash"; + +pub const KNOWN_MODELS: &[&str] = &[ + "claude-opus-4-5@20251101", + "claude-sonnet-4-5@20250929", + "claude-opus-4-1@20250805", + "claude-haiku-4-5@20251001", + "claude-opus-4@20250514", + "claude-sonnet-4@20250514", + "claude-3-5-haiku@20241022", + "claude-3-haiku@20240307", + "gemini-3-pro", + "gemini-3-flash", + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-2.5-flash-lite", + "gemini-2.0-flash", + "gemini-2.0-flash-lite", +]; + /// Represents available GCP Vertex AI models for goose. /// -/// This enum encompasses different model families and their versions -/// that are supported in the GCP Vertex AI platform. +/// This enum encompasses different model families that are supported +/// in the GCP Vertex AI platform. #[derive(Debug, Clone, PartialEq, Eq)] pub enum GcpVertexAIModel { - /// Claude model family with specific versions - Claude(ClaudeVersion), - /// Gemini model family with specific versions - Gemini(GeminiVersion), + /// Claude model family + Claude(String), + /// Gemini model family + Gemini(String), /// MaaS (Model as a Service) models from Model Garden /// Contains (publisher, full_model_name) MaaS(String, String), } -/// Represents available versions of the Claude model for goose. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ClaudeVersion { - /// Claude Sonnet 4 - Sonnet4, - /// Claude Opus 4 - Opus4, - /// Generic Claude model for custom or new versions - Generic(String), -} - -/// Represents available versions of the Gemini model for goose. -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum GeminiVersion { - /// Gemini 1.5 Pro version - Pro15, - /// Gemini 2.0 Flash version - Flash20, - /// Gemini 2.0 Pro Experimental version - Pro20Exp, - /// Gemini 2.5 Pro Experimental version - Pro25Exp, - /// Gemini 2.5 Flash Preview version - Flash25Preview, - /// Gemini 2.5 Pro Preview version - Pro25Preview, - /// Gemini 2.5 Flash version - Flash25, - /// Gemini 2.5 Pro version - Pro25, - /// Generic Gemini model for custom or new versions - Generic(String), -} - impl fmt::Display for GcpVertexAIModel { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let model_id = match self { - Self::Claude(version) => match version { - ClaudeVersion::Sonnet4 => "claude-sonnet-4@20250514", - ClaudeVersion::Opus4 => "claude-opus-4@20250514", - ClaudeVersion::Generic(name) => name, - }, - Self::Gemini(version) => match version { - GeminiVersion::Pro15 => "gemini-1.5-pro-002", - GeminiVersion::Flash20 => "gemini-2.0-flash-001", - GeminiVersion::Pro20Exp => "gemini-2.0-pro-exp-02-05", - GeminiVersion::Pro25Exp => "gemini-2.5-pro-exp-03-25", - GeminiVersion::Flash25Preview => "gemini-2.5-flash-preview-05-20", - GeminiVersion::Pro25Preview => "gemini-2.5-pro-preview-05-06", - GeminiVersion::Flash25 => "gemini-2.5-flash", - GeminiVersion::Pro25 => "gemini-2.5-pro", - GeminiVersion::Generic(name) => name, - }, - Self::MaaS(_, model_name) => model_name, - }; - write!(f, "{model_id}") + match self { + Self::Claude(name) => write!(f, "{name}"), + Self::Gemini(name) => write!(f, "{name}"), + Self::MaaS(_, name) => write!(f, "{name}"), + } } } @@ -164,35 +136,19 @@ impl TryFrom<&str> for GcpVertexAIModel { type Error = ModelError; fn try_from(s: &str) -> Result { - // Known models - match s { - "claude-sonnet-4@20250514" => Ok(Self::Claude(ClaudeVersion::Sonnet4)), - "claude-opus-4@20250514" => Ok(Self::Claude(ClaudeVersion::Opus4)), - "gemini-1.5-pro-002" => Ok(Self::Gemini(GeminiVersion::Pro15)), - "gemini-2.0-flash-001" => Ok(Self::Gemini(GeminiVersion::Flash20)), - "gemini-2.0-pro-exp-02-05" => Ok(Self::Gemini(GeminiVersion::Pro20Exp)), - "gemini-2.5-pro-exp-03-25" => Ok(Self::Gemini(GeminiVersion::Pro25Exp)), - "gemini-2.5-flash-preview-05-20" => Ok(Self::Gemini(GeminiVersion::Flash25Preview)), - "gemini-2.5-pro-preview-05-06" => Ok(Self::Gemini(GeminiVersion::Pro25Preview)), - "gemini-2.5-flash" => Ok(Self::Gemini(GeminiVersion::Flash25)), - "gemini-2.5-pro" => Ok(Self::Gemini(GeminiVersion::Pro25)), - // MaaS models (Model as a Service from Model Garden) - _ if s.ends_with("-maas") => { - let publisher = s - .split('-') - .next() - .ok_or_else(|| ModelError::UnsupportedModel(s.to_string()))? - .to_string(); - Ok(Self::MaaS(publisher, s.to_string())) - } - // Generic models based on prefix matching - _ if s.starts_with("claude-") => { - Ok(Self::Claude(ClaudeVersion::Generic(s.to_string()))) - } - _ if s.starts_with("gemini-") => { - Ok(Self::Gemini(GeminiVersion::Generic(s.to_string()))) - } - _ => Err(ModelError::UnsupportedModel(s.to_string())), + if s.starts_with("claude-") { + Ok(Self::Claude(s.to_string())) + } else if s.starts_with("gemini-") { + Ok(Self::Gemini(s.to_string())) + } else if s.ends_with("-maas") { + let publisher = s + .split('-') + .next() + .ok_or_else(|| ModelError::UnsupportedModel(s.to_string()))? + .to_string(); + Ok(Self::MaaS(publisher, s.to_string())) + } else { + Err(ModelError::UnsupportedModel(s.to_string())) } } } @@ -397,21 +353,16 @@ mod tests { #[test] fn test_model_parsing() -> Result<()> { - let valid_models = [ - "claude-sonnet-4-20250514", - "claude-sonnet-4@20250514", - "gemini-1.5-pro-002", - "gemini-2.0-flash-001", - "gemini-2.0-pro-exp-02-05", - "gemini-2.5-pro-exp-03-25", - "gemini-2.5-flash-preview-05-20", - "gemini-2.5-pro-preview-05-06", - ]; + let claude = GcpVertexAIModel::try_from("claude-sonnet-4@20250514")?; + assert!(matches!(claude, GcpVertexAIModel::Claude(_))); + assert_eq!(claude.to_string(), "claude-sonnet-4@20250514"); - for model_id in valid_models { - let model = GcpVertexAIModel::try_from(model_id)?; - assert_eq!(model.to_string(), model_id); - } + let gemini = GcpVertexAIModel::try_from("gemini-2.5-flash")?; + assert!(matches!(gemini, GcpVertexAIModel::Gemini(_))); + assert_eq!(gemini.to_string(), "gemini-2.5-flash"); + + let maas = GcpVertexAIModel::try_from("qwen-maas")?; + assert!(matches!(maas, GcpVertexAIModel::MaaS(_, _))); assert!(GcpVertexAIModel::try_from("unsupported-model").is_err()); Ok(()) @@ -419,71 +370,24 @@ mod tests { #[test] fn test_default_locations() -> Result<()> { - let test_cases = [ - ("claude-sonnet-4-20250514", GcpLocation::Ohio), - ("claude-sonnet-4@20250514", GcpLocation::Ohio), - ("gemini-1.5-pro-002", GcpLocation::Iowa), - ("gemini-2.0-flash-001", GcpLocation::Iowa), - ("gemini-2.0-pro-exp-02-05", GcpLocation::Iowa), - ("gemini-2.5-pro-exp-03-25", GcpLocation::Iowa), - ("gemini-2.5-flash-preview-05-20", GcpLocation::Iowa), - ("gemini-2.5-pro-preview-05-06", GcpLocation::Iowa), - ]; + let claude_model = GcpVertexAIModel::try_from("claude-sonnet-4@20250514")?; + assert_eq!(claude_model.known_location(), GcpLocation::Ohio); - for (model_id, expected_location) in test_cases { - let model = GcpVertexAIModel::try_from(model_id)?; - assert_eq!( - model.known_location(), - expected_location, - "Model {model_id} should have default location {expected_location:?}", - ); - - let context = RequestContext::new(model_id)?; - assert_eq!( - context.model.known_location(), - expected_location, - "RequestContext for {model_id} should have default location {expected_location:?}", - ); - } + let gemini_model = GcpVertexAIModel::try_from("gemini-2.5-flash")?; + assert_eq!(gemini_model.known_location(), GcpLocation::Iowa); Ok(()) } #[test] - fn test_generic_model_parsing() -> Result<()> { - // Test generic Claude models - let claude_models = [ - "claude-3-8-apex@20250301", - "claude-new-version", - "claude-experimental", - ]; + fn test_unknown_model_parsing() -> Result<()> { + let model = GcpVertexAIModel::try_from("claude-future-version")?; + assert!(matches!(model, GcpVertexAIModel::Claude(_))); + assert_eq!(model.to_string(), "claude-future-version"); - for model_id in claude_models { - let model = GcpVertexAIModel::try_from(model_id)?; - match model { - GcpVertexAIModel::Claude(ClaudeVersion::Generic(ref name)) => { - assert_eq!(name, model_id); - } - _ => panic!("Expected Claude generic model for {model_id}"), - } - assert_eq!(model.to_string(), model_id); - assert_eq!(model.known_location(), GcpLocation::Ohio); - } - - // Test generic Gemini models - let gemini_models = ["gemini-3-pro", "gemini-2.0-flash", "gemini-experimental"]; - - for model_id in gemini_models { - let model = GcpVertexAIModel::try_from(model_id)?; - match model { - GcpVertexAIModel::Gemini(GeminiVersion::Generic(ref name)) => { - assert_eq!(name, model_id); - } - _ => panic!("Expected Gemini generic model for {model_id}"), - } - assert_eq!(model.to_string(), model_id); - assert_eq!(model.known_location(), GcpLocation::Iowa); - } + let model = GcpVertexAIModel::try_from("gemini-4.0-ultra")?; + assert!(matches!(model, GcpVertexAIModel::Gemini(_))); + assert_eq!(model.to_string(), "gemini-4.0-ultra"); Ok(()) } diff --git a/crates/goose/src/providers/gcpvertexai.rs b/crates/goose/src/providers/gcpvertexai.rs index 4fa66055e6..3ae7b1e2d8 100644 --- a/crates/goose/src/providers/gcpvertexai.rs +++ b/crates/goose/src/providers/gcpvertexai.rs @@ -19,11 +19,9 @@ use crate::providers::base::{ConfigKey, MessageStream, Provider, ProviderMetadat use crate::providers::errors::ProviderError; use crate::providers::formats::gcpvertexai::{ - create_request, get_usage, response_to_message, response_to_streaming_message, ClaudeVersion, - GcpVertexAIModel, GeminiVersion, ModelProvider, RequestContext, + create_request, get_usage, response_to_message, response_to_streaming_message, GcpLocation, + ModelProvider, RequestContext, DEFAULT_MODEL, KNOWN_MODELS, }; - -use crate::providers::formats::gcpvertexai::GcpLocation::Iowa; use crate::providers::gcpauth::GcpAuth; use crate::providers::retry::RetryConfig; use crate::providers::utils::RequestLog; @@ -225,7 +223,7 @@ impl GcpVertexAIProvider { .get_param("GCP_LOCATION") .ok() .filter(|location: &String| !location.trim().is_empty()) - .unwrap_or_else(|| Iowa.to_string())) + .unwrap_or_else(|| GcpLocation::Iowa.to_string())) } /// Retrieves an authentication token for API requests. @@ -430,43 +428,118 @@ impl GcpVertexAIProvider { _ => result, } } + + async fn filter_by_org_policy(&self, models: Vec) -> Vec { + let Ok(auth_header) = self.get_auth_header().await else { + tracing::debug!("Could not get auth header for org policy check, returning all models"); + return models; + }; + + let url = format!( + "https://cloudresourcemanager.googleapis.com/v1/projects/{}:getEffectiveOrgPolicy", + self.project_id + ); + + let payload = serde_json::json!({ + "constraint": "constraints/vertexai.allowedModels" + }); + + let response = match self + .client + .post(&url) + .header("Authorization", &auth_header) + .json(&payload) + .send() + .await + { + Ok(r) => r, + Err(e) => { + tracing::debug!("Failed to fetch org policy: {e}, returning all models"); + return models; + } + }; + + let json = match response.json::().await { + Ok(j) => j, + Err(e) => { + tracing::debug!("Failed to parse org policy response: {e}, returning all models"); + return models; + } + }; + + let allowed_patterns: Vec = json + .get("listPolicy") + .and_then(|lp| lp.get("allowedValues")) + .and_then(|av| av.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str()) + .map(|s| s.to_string()) + .collect() + }) + .unwrap_or_default(); + + if allowed_patterns.is_empty() { + return models; + } + + models + .into_iter() + .filter(|model| Self::is_model_allowed(model, &allowed_patterns)) + .collect() + } + + fn is_model_allowed(model: &str, allowed_patterns: &[String]) -> bool { + let publisher = if model.starts_with("claude-") { + "anthropic" + } else if model.starts_with("gemini-") { + "google" + } else { + return true; + }; + + for pattern in allowed_patterns { + if pattern.contains(&format!("publishers/{publisher}/models/*")) { + return true; + } + + let pattern_model = pattern + .split("/models/") + .nth(1) + .map(|s| s.trim_end_matches(":predict").trim_end_matches(":*")); + + if let Some(pattern_model) = pattern_model { + if model == pattern_model || model.starts_with(&format!("{pattern_model}@")) { + return true; + } + } + } + + false + } } #[async_trait] impl Provider for GcpVertexAIProvider { - /// Returns metadata about the GCP Vertex AI provider. fn metadata() -> ProviderMetadata where Self: Sized, { - let model_strings: Vec = [ - GcpVertexAIModel::Claude(ClaudeVersion::Sonnet4), - GcpVertexAIModel::Claude(ClaudeVersion::Opus4), - GcpVertexAIModel::Gemini(GeminiVersion::Pro15), - GcpVertexAIModel::Gemini(GeminiVersion::Flash20), - GcpVertexAIModel::Gemini(GeminiVersion::Pro20Exp), - GcpVertexAIModel::Gemini(GeminiVersion::Pro25Exp), - GcpVertexAIModel::Gemini(GeminiVersion::Flash25Preview), - GcpVertexAIModel::Gemini(GeminiVersion::Pro25Preview), - GcpVertexAIModel::Gemini(GeminiVersion::Flash25), - GcpVertexAIModel::Gemini(GeminiVersion::Pro25), - ] - .iter() - .map(|model| model.to_string()) - .collect(); - - let known_models: Vec<&str> = model_strings.iter().map(|s| s.as_str()).collect(); - ProviderMetadata::new( "gcp_vertex_ai", "GCP Vertex AI", "Access variety of AI models such as Claude, Gemini through Vertex AI", - "gemini-2.5-flash", - known_models, + DEFAULT_MODEL, + KNOWN_MODELS.to_vec(), GCP_VERTEX_AI_DOC_URL, vec![ ConfigKey::new("GCP_PROJECT_ID", true, false, None), - ConfigKey::new("GCP_LOCATION", true, false, Some(Iowa.to_string().as_str())), + ConfigKey::new( + "GCP_LOCATION", + true, + false, + Some(&GcpLocation::Iowa.to_string()), + ), ConfigKey::new( "GCP_MAX_RETRIES", false, @@ -493,6 +566,7 @@ impl Provider for GcpVertexAIProvider { ), ], ) + .with_unlisted_models() } fn get_name(&self) -> &str { @@ -587,6 +661,12 @@ impl Provider for GcpVertexAIProvider { } })) } + + async fn fetch_supported_models(&self) -> Result>, ProviderError> { + let models: Vec = KNOWN_MODELS.iter().map(|s| s.to_string()).collect(); + let filtered = self.filter_by_org_policy(models).await; + Ok(Some(filtered)) + } } #[cfg(test)] @@ -705,15 +785,9 @@ mod tests { #[test] fn test_provider_metadata() { let metadata = GcpVertexAIProvider::metadata(); - let model_names: Vec = metadata - .known_models - .iter() - .map(|m| m.name.clone()) - .collect(); - assert!(model_names.contains(&"claude-sonnet-4@20250514".to_string())); - assert!(model_names.contains(&"gemini-1.5-pro-002".to_string())); - assert!(model_names.contains(&"gemini-2.5-pro".to_string())); - // Should contain the original 2 config keys plus 4 new retry-related ones + assert!(!metadata.known_models.is_empty()); + assert_eq!(metadata.default_model, "gemini-2.5-flash"); assert_eq!(metadata.config_keys.len(), 6); + assert!(metadata.allows_unlisted_models); } } diff --git a/crates/goose/src/providers/provider_registry.rs b/crates/goose/src/providers/provider_registry.rs index a89207a265..f10df127d3 100644 --- a/crates/goose/src/providers/provider_registry.rs +++ b/crates/goose/src/providers/provider_registry.rs @@ -114,6 +114,7 @@ impl ProviderRegistry { known_models, model_doc_link: base_metadata.model_doc_link, config_keys, + allows_unlisted_models: false, }; self.entries.insert( diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 96c156fff3..9ba0a903c4 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -397,6 +397,7 @@ mod tests { known_models: vec![], model_doc_link: "".to_string(), config_keys: vec![], + allows_unlisted_models: false, } } diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 6961c28044..0f132afd64 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -4735,6 +4735,10 @@ "config_keys" ], "properties": { + "allows_unlisted_models": { + "type": "boolean", + "description": "Whether this provider allows entering model names not in the fetched list" + }, "config_keys": { "type": "array", "items": { diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index a7a12e9468..b253d957f6 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -630,6 +630,10 @@ export type ProviderEngine = 'openai' | 'ollama' | 'anthropic'; * Metadata about a provider's configuration requirements and capabilities */ export type ProviderMetadata = { + /** + * Whether this provider allows entering model names not in the fetched list + */ + allows_unlisted_models?: boolean; /** * Required configuration keys */ diff --git a/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx b/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx index e58932ea53..d65a8ab35f 100644 --- a/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx +++ b/ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx @@ -119,12 +119,17 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps) }); }); } + // 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, + }); + } }); } - // Append a simple "custom" option to enable free-text entry - options.push({ value: '__custom__', label: 'Use custom model…', provider: '' }); - setModelOptions(options); } catch (error) { console.error('Error loading configuration:', error); @@ -241,9 +246,10 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps) onChange={(newValue: unknown) => { const option = newValue as { value: string; provider: string } | null; if (option) { - if (option.value === '__custom__') { + if (option.value.startsWith('__custom__')) { setIsLeadCustomModel(true); setLeadModel(''); + setLeadProvider(option.provider); return; } setLeadModel(option.value); @@ -294,9 +300,10 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps) onChange={(newValue: unknown) => { const option = newValue as { value: string; provider: string } | null; if (option) { - if (option.value === '__custom__') { + if (option.value.startsWith('__custom__')) { setIsWorkerCustomModel(true); setWorkerModel(''); + setWorkerProvider(option.provider); return; } setWorkerModel(option.value); diff --git a/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx b/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx index f6d2f485e0..cd9c7feff6 100644 --- a/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx +++ b/ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx @@ -213,29 +213,34 @@ export const SwitchModelModal = ({ const errors: string[] = []; results.forEach(({ provider: p, models, error }) => { + const modelList = error + ? (p.metadata.known_models?.map(({ name }) => name) || []) + : (models || []); + if (error) { errors.push(error); - // Fallback to metadata known_models on error - if (p.metadata.known_models && p.metadata.known_models.length > 0) { - groupedOptions.push({ - options: p.metadata.known_models.map(({ name }) => ({ - value: name, - label: name, - providerType: p.provider_type, - provider: p.name, - })), - }); - } - } else if (models && models.length > 0) { - groupedOptions.push({ - options: models.map((m) => ({ - value: m, - label: m, - provider: p.name, - providerType: p.provider_type, - })), + } + + const options: { value: string; label: string; provider: string; providerType: ProviderType }[] = + modelList.map((m) => ({ + value: m, + label: m, + provider: p.name, + providerType: p.provider_type, + })); + + if (p.metadata.allows_unlisted_models && p.provider_type !== 'Custom') { + options.push({ + value: 'custom', + label: 'Enter a model not listed...', + provider: p.name, + providerType: p.provider_type, }); } + + if (options.length > 0) { + groupedOptions.push({ options }); + } }); // Log errors if any providers failed (don't show to user) @@ -243,20 +248,6 @@ export const SwitchModelModal = ({ console.error('Provider model fetch errors:', errors); } - // Add the "Custom model" option to each provider group - groupedOptions.forEach((group) => { - const option = group.options[0]; - const providerName = option?.provider; - if (providerName && option?.providerType !== 'Custom') { - group.options.push({ - value: 'custom', - label: 'Use custom model', - provider: providerName, - providerType: option?.providerType, - }); - } - }); - setModelOptions(groupedOptions); setOriginalModelOptions(groupedOptions); } catch (error: unknown) { @@ -293,6 +284,7 @@ export const SwitchModelModal = ({ if (selectedOption?.value === 'custom') { setIsCustomModel(true); setModel(''); + setProvider(selectedOption.provider); setUserClearedModel(false); } else if (selectedOption === null) { // User cleared the selection @@ -302,6 +294,7 @@ export const SwitchModelModal = ({ } else { setIsCustomModel(false); setModel(selectedOption?.value || ''); + setProvider(selectedOption?.provider || ''); setUserClearedModel(false); } };