From 176ab18880fabffdc636e25fa0ea424779e458af Mon Sep 17 00:00:00 2001
From: xinbenlv <zzn-github@zzn.im>
Date: Wed, 4 Jun 2025 10:58:30 -0700
Subject: [PATCH] feat(providers): Add support for Gemini 2.5 Flash Preview and
 Pro Preview models (#2780)

---
 .../goose-server/src/routes/providers_and_keys.json  |  2 +-
 crates/goose/src/providers/formats/gcpvertexai.rs    | 12 ++++++++++++
 crates/goose/src/providers/gcpvertexai.rs            |  3 +++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/crates/goose-server/src/routes/providers_and_keys.json b/crates/goose-server/src/routes/providers_and_keys.json
index 830cf665..fc61c5b9 100644
--- a/crates/goose-server/src/routes/providers_and_keys.json
+++ b/crates/goose-server/src/routes/providers_and_keys.json
@@ -20,7 +20,7 @@
     "gcp_vertex_ai": {
         "name": "GCP Vertex AI",
         "description": "Use Vertex AI platform models",
-        "models": ["claude-3-5-haiku@20241022", "claude-3-5-sonnet@20240620", "claude-3-5-sonnet-v2@20241022", "claude-3-7-sonnet@20250219", "gemini-1.5-pro-002", "gemini-2.0-flash-001", "gemini-2.0-pro-exp-02-05", "gemini-2.5-pro-exp-03-25"],
+        "models": ["claude-3-5-haiku@20241022", "claude-3-5-sonnet@20240620", "claude-3-5-sonnet-v2@20241022", "claude-3-7-sonnet@20250219", "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"],
         "required_keys": ["GCP_PROJECT_ID", "GCP_LOCATION"]
     },
     "google": {
diff --git a/crates/goose/src/providers/formats/gcpvertexai.rs b/crates/goose/src/providers/formats/gcpvertexai.rs
index 5bc94841..d83d1939 100644
--- a/crates/goose/src/providers/formats/gcpvertexai.rs
+++ b/crates/goose/src/providers/formats/gcpvertexai.rs
@@ -98,6 +98,10 @@ pub enum GeminiVersion {
     Pro20Exp,
     /// Gemini 2.5 Pro Experimental version
     Pro25Exp,
+    /// Gemini 2.5 Flash Preview version
+    Flash25Preview,
+    /// Gemini 2.5 Pro Preview version
+    Pro25Preview,
     /// Generic Gemini model for custom or new versions
     Generic(String),
 }
@@ -118,6 +122,8 @@ impl fmt::Display for GcpVertexAIModel {
                 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::Generic(name) => name,
             },
         };
@@ -154,6 +160,8 @@ impl TryFrom<&str> for GcpVertexAIModel {
             "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)),
             // Generic models based on prefix matching
             _ if s.starts_with("claude-") => {
                 Ok(Self::Claude(ClaudeVersion::Generic(s.to_string())))
@@ -349,6 +357,8 @@ mod tests {
             "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",
         ];
 
         for model_id in valid_models {
@@ -372,6 +382,8 @@ mod tests {
             ("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),
         ];
 
         for (model_id, expected_location) in test_cases {
diff --git a/crates/goose/src/providers/gcpvertexai.rs b/crates/goose/src/providers/gcpvertexai.rs
index afec862d..6385ec29 100644
--- a/crates/goose/src/providers/gcpvertexai.rs
+++ b/crates/goose/src/providers/gcpvertexai.rs
@@ -434,6 +434,9 @@ impl Provider for GcpVertexAIProvider {
             GcpVertexAIModel::Gemini(GeminiVersion::Pro15),
             GcpVertexAIModel::Gemini(GeminiVersion::Flash20),
             GcpVertexAIModel::Gemini(GeminiVersion::Pro20Exp),
+            GcpVertexAIModel::Gemini(GeminiVersion::Pro25Exp),
+            GcpVertexAIModel::Gemini(GeminiVersion::Flash25Preview),
+            GcpVertexAIModel::Gemini(GeminiVersion::Pro25Preview),
         ]
         .iter()
         .map(|model| model.to_string())