From ac6524cdb9fdb15391a07437db5476b1a18bb3b4 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Fri, 1 May 2026 11:18:35 -0500 Subject: [PATCH] Fix Vertex Gemma 4 routing metadata --- litellm/llms/vertex_ai/common_utils.py | 7 ++-- ...odel_prices_and_context_window_backup.json | 7 ++-- model_prices_and_context_window.json | 7 ++-- .../test_vertex_gemma_transformation.py | 35 +++++++++++++++++++ 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index b4bfde5f541..43f4edd19b7 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -50,7 +50,7 @@ def get_vertex_ai_model_route( Determine which handler to use for a Vertex AI model based on the model name. Args: - model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "openai/gpt-oss-120b") + model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "google/gemma-4-26b-a4b-it-maas", "openai/gpt-oss-120b") litellm_params: Optional litellm parameters dict that may contain base_model for routing Returns: @@ -66,6 +66,9 @@ def get_vertex_ai_model_route( >>> get_vertex_ai_model_route("gemma/gemma-3-12b-it") VertexAIModelRoute.GEMMA + >>> get_vertex_ai_model_route("google/gemma-4-26b-a4b-it-maas") + VertexAIModelRoute.GEMMA + >>> get_vertex_ai_model_route("openai/gpt-oss-120b") VertexAIModelRoute.MODEL_GARDEN @@ -99,7 +102,7 @@ def get_vertex_ai_model_route( return VertexAIModelRoute.BGE # Check for gemma models - if "gemma/" in model: + if "gemma/" in model or model.startswith("google/gemma-"): return VertexAIModelRoute.GEMMA # Check for model garden openai models diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f758128dcb6..405e41ef3e2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32735,14 +32735,15 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 256000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 262144, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supports_function_calling": true, "supports_system_messages": true, + "supports_tool_choice": true, "supports_vision": true }, "vertex_ai/deep-research-pro-preview-12-2025": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a0b43678d9a..6a8e6f11f18 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32789,14 +32789,15 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai", - "max_input_tokens": 256000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_input_tokens": 262144, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supports_function_calling": true, "supports_system_messages": true, + "supports_tool_choice": true, "supports_vision": true }, "vertex_ai/deep-research-pro-preview-12-2025": { diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 3e3e8901706..4e1b07c3d22 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -5,6 +5,7 @@ Maps to: litellm/llms/vertex_ai/vertex_gemma_models/transformation.py """ import json +from pathlib import Path from unittest.mock import AsyncMock, Mock, patch import pytest @@ -20,6 +21,40 @@ def _reset_litellm_http_client_cache(): in_memory_llm_clients_cache.flush_cache() +def test_gemma_4_26b_pricing_metadata_uses_vertex_gemma_route(): + from litellm.llms.vertex_ai.common_utils import ( + VertexAIModelRoute, + get_vertex_ai_model_route, + ) + + model = "google/gemma-4-26b-a4b-it-maas" + litellm_model = f"vertex_ai/{model}" + + assert get_vertex_ai_model_route(model=model) == VertexAIModelRoute.GEMMA + assert ( + get_vertex_ai_model_route(model="gemma/gemma-3-12b-it") + == VertexAIModelRoute.GEMMA + ) + + repo_root = Path(litellm.__file__).resolve().parents[1] + model_cost_files = [ + repo_root / "model_prices_and_context_window.json", + repo_root / "litellm" / "model_prices_and_context_window_backup.json", + ] + + for model_cost_file in model_cost_files: + model_cost = json.loads(model_cost_file.read_text(encoding="utf-8")) + + assert "vertex_ai/gemma/gemma-4-26b-a4b-it-maas" not in model_cost + model_info = model_cost[litellm_model] + + assert model_info["max_input_tokens"] == 262144 + assert model_info["max_output_tokens"] == 128000 + assert model_info["max_tokens"] == 128000 + assert model_info["supports_function_calling"] is True + assert model_info["supports_tool_choice"] is True + + class TestVertexGemmaCompletion: """Test completion flow for Vertex AI Gemma models using litellm.acompletion()"""