tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat(gcp-vertex): add model list with org policy filtering (#6393)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-01-20 22:19:10 +05:30
committed by GitHub
parent a48fd4dc74
commit fbe46d0b22
10 changed files with 283 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/configure.rs
+48 -12
View File
@@ -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),
crates/goose/src/providers/base.rs
+12
View File
@@ -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
crates/goose/src/providers/formats/gcpvertexai.rs
+65 -161
View File
@@ -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(())
}
crates/goose/src/providers/gcpvertexai.rs
+110 -36
View File
@@ -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);
}
}
crates/goose/src/providers/provider_registry.rs
+1
View File
@@ -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(
crates/goose/tests/agent.rs
+1
View File
@@ -397,6 +397,7 @@ mod tests {
known_models: vec![],
model_doc_link: "".to_string(),
config_keys: vec![],
allows_unlisted_models: false,
}
}