From 29e285cba986c576f9a38e36b1af780bd7d93e10 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Fri, 1 May 2026 14:37:47 -0500 Subject: [PATCH] Route Vertex Gemma MaaS through model garden --- litellm/llms/vertex_ai/common_utils.py | 11 +++-- .../test_vertex_model_garden_openapi.py | 47 ++++++++++++++++++- .../test_vertex_gemma_transformation.py | 35 -------------- 3 files changed, 54 insertions(+), 39 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index bfffadd7aa7..38e6784cca0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -114,7 +114,7 @@ def get_vertex_ai_model_route( VertexAIModelRoute.GEMMA >>> get_vertex_ai_model_route("google/gemma-4-26b-a4b-it-maas") - VertexAIModelRoute.GEMMA + VertexAIModelRoute.MODEL_GARDEN >>> get_vertex_ai_model_route("openai/gpt-oss-120b") VertexAIModelRoute.MODEL_GARDEN @@ -152,14 +152,19 @@ def get_vertex_ai_model_route( return VertexAIModelRoute.BGE # Check for gemma models - if "gemma/" in model or model.startswith("google/gemma-"): + if "gemma/" in model: return VertexAIModelRoute.GEMMA # Check for model garden OpenAI-compatible publisher models. # Examples: + # - google/gemma-4-26b-a4b-it-maas # - openai/gpt-oss-120b-maas # - xai/grok-4.1-fast-non-reasoning - if "openai" in model or model.startswith("xai/"): + if ( + "openai" in model + or model.startswith("google/gemma-") + or model.startswith("xai/") + ): return VertexAIModelRoute.MODEL_GARDEN # Check for gemini models diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py index 91261b63252..f99df2d3626 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py @@ -1,7 +1,15 @@ """Vertex Model Garden: OpenAPI base URL for publisher/model ids vs per-endpoint path.""" +import json +from pathlib import Path + import pytest +import litellm +from litellm.llms.vertex_ai.common_utils import ( + VertexAIModelRoute, + get_vertex_ai_model_route, +) from litellm.llms.vertex_ai.vertex_model_garden.main import ( _vertex_model_garden_model_id_in_json_body, create_vertex_url, @@ -11,6 +19,7 @@ from litellm.llms.vertex_ai.vertex_model_garden.main import ( @pytest.mark.parametrize( "model,expect_openapi_base", [ + ("google/gemma-4-26b-a4b-it-maas", True), ("xai/grok-4.1-fast-reasoning", True), ("openai/foo/bar", True), ("5464397967697903616", False), @@ -37,5 +46,41 @@ def test_create_vertex_url_openapi_vs_deployed_endpoint( def test_model_id_in_json_body_heuristic() -> None: - assert _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") is True + assert ( + _vertex_model_garden_model_id_in_json_body("google/gemma-4-26b-a4b-it-maas") + is True + ) + assert ( + _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") + is True + ) assert _vertex_model_garden_model_id_in_json_body("5464397967697903616") is False + + +def test_gemma_4_26b_pricing_metadata_uses_vertex_model_garden_route(): + model = "google/gemma-4-26b-a4b-it-maas" + litellm_model = f"vertex_ai/{model}" + + assert get_vertex_ai_model_route(model=model) == VertexAIModelRoute.MODEL_GARDEN + 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 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 4e1b07c3d22..3e3e8901706 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,7 +5,6 @@ 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 @@ -21,40 +20,6 @@ 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()"""