mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
feat(gcp-vertex): add model list with org policy filtering (#6393)
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
@@ -436,12 +436,12 @@ fn select_model_from_list(
|
||||
provider_meta: &goose::providers::base::ProviderMetadata,
|
||||
) -> anyhow::Result<String> {
|
||||
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<String> = 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::<Vec<_>>(),
|
||||
)
|
||||
.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<String> {
|
||||
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<bool> {
|
||||
match config.set_secret(key_name, &value) {
|
||||
Ok(_) => Ok(true),
|
||||
|
||||
@@ -108,6 +108,9 @@ pub struct ProviderMetadata {
|
||||
pub model_doc_link: String,
|
||||
/// Required configuration keys
|
||||
pub config_keys: Vec<ConfigKey>,
|
||||
/// 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
|
||||
|
||||
@@ -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<Self, Self::Error> {
|
||||
// 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(())
|
||||
}
|
||||
|
||||
@@ -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<String>) -> Vec<String> {
|
||||
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::<Value>().await {
|
||||
Ok(j) => j,
|
||||
Err(e) => {
|
||||
tracing::debug!("Failed to parse org policy response: {e}, returning all models");
|
||||
return models;
|
||||
}
|
||||
};
|
||||
|
||||
let allowed_patterns: Vec<String> = 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<String> = [
|
||||
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<Option<Vec<String>>, ProviderError> {
|
||||
let models: Vec<String> = 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<String> = 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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -397,6 +397,7 @@ mod tests {
|
||||
known_models: vec![],
|
||||
model_doc_link: "".to_string(),
|
||||
config_keys: vec![],
|
||||
allows_unlisted_models: false,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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
|
||||
*/
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user