From 5bf555077f5518267185a71270fa958ffaf3d1c5 Mon Sep 17 00:00:00 2001
From: pedrohsdb <pedro@skyvern.com>
Date: Fri, 27 Feb 2026 19:13:23 -0800
Subject: [PATCH] Consolidate Gemini 3 Pro into version-agnostic
 VERTEX_GEMINI_3_PRO (#4926)

---
 skyvern/config.py                            |  3 +-
 skyvern/forge/sdk/api/llm/config_registry.py | 33 ++++++++++----------
 tests/unit/test_text_prompt_block.py         |  2 +-
 3 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/skyvern/config.py b/skyvern/config.py
index c6eb29295..4bc47539a 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -471,9 +471,8 @@ class Settings(BaseSettings):
                 "llm_key": "VERTEX_GEMINI_2.5_FLASH",
                 "label": "Gemini 2.5 Flash",
             },
-            "gemini-3-pro-preview": {"llm_key": "VERTEX_GEMINI_3.0_PRO", "label": "Gemini 3 Pro"},
+            "gemini-3-pro-preview": {"llm_key": "VERTEX_GEMINI_3_PRO", "label": "Gemini 3 Pro (Latest)"},
             "gemini-3.0-flash": {"llm_key": "VERTEX_GEMINI_3.0_FLASH", "label": "Gemini 3 Flash"},
-            "gemini-3.1-pro-preview": {"llm_key": "VERTEX_GEMINI_3.1_PRO", "label": "Gemini 3.1 Pro"},
             "mercury-2": {"llm_key": "INCEPTION_MERCURY_2", "label": "Inception Mercury 2"},
             "gemini-2.5-flash-lite": {
                 "llm_key": "VERTEX_GEMINI_2.5_FLASH_LITE",
diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py
index 737f490e4..b02b99c32 100644
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -1034,7 +1034,7 @@ if settings.ENABLE_GEMINI:
         ),
     )
     LLMConfigRegistry.register_config(
-        "GEMINI_3.1_PRO",
+        "GEMINI_3_PRO",
         LLMConfig(
             "gemini/gemini-3.1-pro-preview",
             ["GEMINI_API_KEY"],
@@ -1046,6 +1046,11 @@ if settings.ENABLE_GEMINI:
             ),
         ),
     )
+    # Backward compat alias for non-Vertex Gemini 3 Pro
+    LLMConfigRegistry.register_config(
+        "GEMINI_3.1_PRO",
+        LLMConfigRegistry.get_config("GEMINI_3_PRO"),
+    )
 
 
 if settings.ENABLE_NOVITA:
@@ -1381,15 +1386,15 @@ if settings.ENABLE_VERTEX_AI:
         ),
     )
     LLMConfigRegistry.register_config(
-        "VERTEX_GEMINI_3.0_PRO",
+        "VERTEX_GEMINI_3_PRO",
         LLMConfig(
-            "vertex_ai/gemini-3-pro-preview",
+            "vertex_ai/gemini-3.1-pro-preview",
             [],
             supports_vision=True,
             add_assistant_prefix=False,
             max_completion_tokens=65536,
             litellm_params=LiteLLMParams(
-                api_base=f"{api_base}/gemini-3-pro-preview" if api_base else None,
+                api_base=f"{api_base}/gemini-3.1-pro-preview" if api_base else None,
                 vertex_location=settings.VERTEX_LOCATION,
                 thinking_level="medium" if settings.GEMINI_INCLUDE_THOUGHT else "minimal",
                 vertex_credentials=settings.VERTEX_CREDENTIALS,
@@ -1412,21 +1417,15 @@ if settings.ENABLE_VERTEX_AI:
             ),
         ),
     )
+    # Backward compat aliases — both resolve to the canonical VERTEX_GEMINI_3_PRO.
+    # Bump VERTEX_GEMINI_3_PRO above when Google ships a newer version.
+    LLMConfigRegistry.register_config(
+        "VERTEX_GEMINI_3.0_PRO",
+        LLMConfigRegistry.get_config("VERTEX_GEMINI_3_PRO"),
+    )
     LLMConfigRegistry.register_config(
         "VERTEX_GEMINI_3.1_PRO",
-        LLMConfig(
-            "vertex_ai/gemini-3.1-pro-preview",
-            [],
-            supports_vision=True,
-            add_assistant_prefix=False,
-            max_completion_tokens=65536,
-            litellm_params=LiteLLMParams(
-                api_base=f"{api_base}/gemini-3.1-pro-preview" if api_base else None,
-                vertex_location=settings.VERTEX_LOCATION,
-                thinking_level="medium" if settings.GEMINI_INCLUDE_THOUGHT else "minimal",
-                vertex_credentials=settings.VERTEX_CREDENTIALS,
-            ),
-        ),
+        LLMConfigRegistry.get_config("VERTEX_GEMINI_3_PRO"),
     )
     LLMConfigRegistry.register_config(
         "VERTEX_GEMINI_2.5_FLASH_LITE",
diff --git a/tests/unit/test_text_prompt_block.py b/tests/unit/test_text_prompt_block.py
index 2f8026951..18bed8ff5 100644
--- a/tests/unit/test_text_prompt_block.py
+++ b/tests/unit/test_text_prompt_block.py
@@ -18,7 +18,7 @@ block_module = sys.modules["skyvern.forge.sdk.workflow.models.block"]
     ("model_name", "expected_llm_key"),
     [
         ("gemini-2.5-flash", "VERTEX_GEMINI_2.5_FLASH"),
-        ("gemini-3-pro-preview", "VERTEX_GEMINI_3.0_PRO"),
+        ("gemini-3-pro-preview", "VERTEX_GEMINI_3_PRO"),
     ],
 )
 async def test_text_prompt_block_uses_selected_model(monkeypatch, model_name, expected_llm_key):