scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path

Browse source 
Fixes #26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on #26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.
This commit is contained in:
Iana 2026-05-15 18:55:38 +01:00
parent ed0138b50e
commit e3da87f739
5 changed files with 378 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -41,6 +41,7 @@ class PartnerModelPrefixes(str, Enum):
MINIMAX_PREFIX = "minimaxai/"
MOONSHOT_PREFIX = "moonshotai/"
ZAI_PREFIX = "zai-org/"
GEMMA_MAAS_PREFIX = "google/gemma-"
class VertexAIPartnerModels(VertexBase):
@ -68,6 +69,7 @@ class VertexAIPartnerModels(VertexBase):
or model.startswith(PartnerModelPrefixes.MINIMAX_PREFIX)
or model.startswith(PartnerModelPrefixes.MOONSHOT_PREFIX)
or model.startswith(PartnerModelPrefixes.ZAI_PREFIX)
or model.startswith(PartnerModelPrefixes.GEMMA_MAAS_PREFIX)
):
return True
return False
@ -82,6 +84,7 @@ class VertexAIPartnerModels(VertexBase):
PartnerModelPrefixes.MINIMAX_PREFIX,
PartnerModelPrefixes.MOONSHOT_PREFIX,
PartnerModelPrefixes.ZAI_PREFIX,
PartnerModelPrefixes.GEMMA_MAAS_PREFIX,
]
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
return True

16
model_prices_and_context_window.json
View file

@ -31996,6 +31996,22 @@
"us-central1"
]
},
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 256000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/maas/google/gemma-4-26b-a4b-it",
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"vertex_ai/openai/gpt-oss-120b-maas": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "vertex_ai-openai_models",

258
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -440,7 +440,9 @@ def test_vertex_ai_complex_response_schema():
optional_params = {}
v.apply_response_schema_transformation(
value=non_default_params["response_format"], optional_params=optional_params, model="gemini-1.5-pro-preview-0409"
value=non_default_params["response_format"],
optional_params=optional_params,
model="gemini-1.5-pro-preview-0409",
)
# Assertions for the transformed schema
@ -558,7 +560,6 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
assert url == expected_url
@pytest.mark.parametrize(
"model_cost_entry, vertex_region, expected_region",
[
@ -571,9 +572,17 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
# Model with supported_regions=["us-west2"], no user region -> use "us-west2"
({"supported_regions": ["us-west2"]}, None, "us-west2"),
# Model with supported_regions=["us-west2", "us-central1"], user passes supported region -> respect it
({"supported_regions": ["us-west2", "us-central1"]}, "us-central1", "us-central1"),
(
{"supported_regions": ["us-west2", "us-central1"]},
"us-central1",
"us-central1",
),
# Model with supported_regions=["us-west2", "us-central1"], user passes unsupported region -> override
({"supported_regions": ["us-west2", "us-central1"]}, "europe-west1", "us-west2"),
(
{"supported_regions": ["us-west2", "us-central1"]},
"europe-west1",
"us-west2",
),
# No model_cost entry, no user region -> default us-central1
({}, None, "us-central1"),
# No model_cost entry, user specifies region -> use specified region
@ -656,11 +665,12 @@ def test_vertex_filter_format_uri():
assert "uri" not in json.dumps(new_parameters)
def test_convert_schema_types_type_array_conversion():
"""
Test _convert_schema_types function handles type arrays and case conversion.
This test verifies the fix for the issue where type arrays like ["string", "number"]
This test verifies the fix for the issue where type arrays like ["string", "number"]
would raise an exception in Vertex AI schema validation.
Relevant issue: https://github.com/BerriAI/litellm/issues/14091
@ -673,12 +683,12 @@ def test_convert_schema_types_type_array_conversion():
"properties": {
"studio": {
"type": ["string", "number"],
"description": "The studio ID or name"
"description": "The studio ID or name",
}
},
"required": ["studio"],
"additionalProperties": False,
"$schema": "http://json-schema.org/draft-07/schema#"
"$schema": "http://json-schema.org/draft-07/schema#",
}
# Expected output: Vertex AI compatible schema with anyOf and uppercase types
@ -686,16 +696,13 @@ def test_convert_schema_types_type_array_conversion():
"type": "object",
"properties": {
"studio": {
"anyOf": [
{"type": "string"},
{"type": "number"}
],
"description": "The studio ID or name"
"anyOf": [{"type": "string"}, {"type": "number"}],
"description": "The studio ID or name",
}
},
"required": ["studio"],
"additionalProperties": False,
"$schema": "http://json-schema.org/draft-07/schema#"
"$schema": "http://json-schema.org/draft-07/schema#",
}
# Apply the transformation
@ -718,15 +725,17 @@ def test_convert_schema_types_type_array_conversion():
assert anyof_types[1]["type"] == "number"
# 4. Other properties preserved
assert input_schema["properties"]["studio"]["description"] == "The studio ID or name"
assert (
input_schema["properties"]["studio"]["description"] == "The studio ID or name"
)
assert input_schema["required"] == ["studio"]
def test_fix_enum_empty_strings():
"""
Test _fix_enum_empty_strings function replaces empty strings with None in enum arrays.
This test verifies the fix for the issue where Gemini rejects tool definitions
This test verifies the fix for the issue where Gemini rejects tool definitions
with empty strings in enum values, causing API failures.
Relevant issue: Gemini does not accept empty strings in enum values
@ -740,23 +749,23 @@ def test_fix_enum_empty_strings():
"user_agent_type": {
"enum": ["", "desktop", "mobile", "tablet"],
"type": "string",
"description": "Device type for user agent"
"description": "Device type for user agent",
}
},
"required": ["user_agent_type"]
"required": ["user_agent_type"],
}
# Expected output: Empty strings replaced with None
expected_output = {
"type": "object",
"type": "object",
"properties": {
"user_agent_type": {
"enum": [None, "desktop", "mobile", "tablet"],
"type": "string",
"description": "Device type for user agent"
"description": "Device type for user agent",
}
},
"required": ["user_agent_type"]
"required": ["user_agent_type"],
}
# Apply the transformation
@ -859,7 +868,7 @@ def test_construct_target_url_with_version_prefix():
def test_fix_enum_types():
"""
Test _fix_enum_types function removes enum fields when type is not string.
This test verifies the fix for the issue where Gemini rejects cached content
with function parameter enums on non-string types, causing API failures.
@ -874,38 +883,41 @@ def test_fix_enum_types():
"truncateMode": {
"enum": ["auto", "none", "start", "end"],
"type": "string", # This should keep the enum
"description": "How to truncate content"
"description": "How to truncate content",
},
"maxLength": {
"enum": [100, 200, 500], # This should be removed
"type": "integer",
"description": "Maximum length"
"description": "Maximum length",
},
"enabled": {
"enum": [True, False], # This should be removed
"type": "boolean",
"description": "Whether feature is enabled"
"description": "Whether feature is enabled",
},
"nested": {
"type": "object",
"properties": {
"innerEnum": {
"enum": ["a", "b", "c"], # This should be kept
"type": "string"
"type": "string",
},
"innerNonStringEnum": {
"enum": [1, 2, 3], # This should be removed
"type": "integer"
}
}
"type": "integer",
},
},
},
"anyOfField": {
"anyOf": [
{"type": "string", "enum": ["option1", "option2"]}, # This should be kept
{"type": "integer", "enum": [1, 2, 3]} # This should be removed
{
"type": "string",
"enum": ["option1", "option2"],
}, # This should be kept
{"type": "integer", "enum": [1, 2, 3]}, # This should be removed
]
}
}
},
},
}
# Expected output: Non-string enums removed, string enums kept
@ -919,31 +931,32 @@ def test_fix_enum_types():
},
"maxLength": { # enum removed
"type": "integer",
"description": "Maximum length"
"description": "Maximum length",
},
"enabled": { # enum removed
"type": "boolean",
"description": "Whether feature is enabled"
"description": "Whether feature is enabled",
},
"nested": {
"type": "object",
"properties": {
"innerEnum": {
"enum": ["a", "b", "c"], # Kept - string type
"type": "string"
"type": "string",
},
"innerNonStringEnum": { # enum removed
"type": "integer"
}
}
"innerNonStringEnum": {"type": "integer"}, # enum removed
},
},
"anyOfField": {
"anyOf": [
{"type": "string", "enum": ["option1", "option2"]}, # Kept - has string type
{"type": "integer"} # enum removed
{
"type": "string",
"enum": ["option1", "option2"],
}, # Kept - has string type
{"type": "integer"}, # enum removed
]
}
}
},
},
}
# Apply the transformation
@ -955,15 +968,27 @@ def test_fix_enum_types():
# Verify specific transformations:
# 1. String enums are preserved
assert "enum" in input_schema["properties"]["truncateMode"]
assert input_schema["properties"]["truncateMode"]["enum"] == ["auto", "none", "start", "end"]
assert input_schema["properties"]["truncateMode"]["enum"] == [
"auto",
"none",
"start",
"end",
]
assert "enum" in input_schema["properties"]["nested"]["properties"]["innerEnum"]
assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == ["a", "b", "c"]
assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == [
"a",
"b",
"c",
]
# 2. Non-string enums are removed
assert "enum" not in input_schema["properties"]["maxLength"]
assert "enum" not in input_schema["properties"]["enabled"]
assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
assert (
"enum"
not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
)
# 3. anyOf with string type keeps enum, non-string removes it
assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0]
@ -1003,8 +1028,6 @@ def test_get_token_url():
print("url=", url)
should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features(
optional_params={"temperature": 0.1}
)
@ -1210,9 +1233,7 @@ def test_vertex_ai_minimax_uses_openai_handler():
VertexAIPartnerModels,
)
assert VertexAIPartnerModels.should_use_openai_handler(
"minimaxai/minimax-m2-maas"
)
assert VertexAIPartnerModels.should_use_openai_handler("minimaxai/minimax-m2-maas")
def test_vertex_ai_moonshot_uses_openai_handler():
@ -1236,9 +1257,7 @@ def test_vertex_ai_zai_uses_openai_handler():
VertexAIPartnerModels,
)
assert VertexAIPartnerModels.should_use_openai_handler(
"zai-org/glm-4.7-maas"
)
assert VertexAIPartnerModels.should_use_openai_handler("zai-org/glm-4.7-maas")
def test_vertex_ai_zai_is_partner_model():
@ -1252,17 +1271,74 @@ def test_vertex_ai_zai_is_partner_model():
assert VertexAIPartnerModels.is_vertex_partner_model("zai-org/glm-4.7-maas")
def test_vertex_ai_gemma_maas_is_partner_model():
"""
Ensure Gemma MaaS models are detected as Vertex AI partner models so they
route through the OpenAI-compatible /endpoints/openapi path (not the
legacy non-gemini path or the vertex_ai/gemma/ predict-endpoint handler).
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
)
assert VertexAIPartnerModels.is_vertex_partner_model(
"google/gemma-4-26b-a4b-it-maas"
)
def test_vertex_ai_gemma_maas_uses_openai_handler():
"""
Ensure Gemma MaaS partner models re-use the OpenAI-format handler.
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
)
assert VertexAIPartnerModels.should_use_openai_handler(
"google/gemma-4-26b-a4b-it-maas"
)
def test_vertex_ai_gemma_maas_routes_to_partner_models():
"""
Regression guard for owtaylor's worry that Gemma MaaS could be misrouted as
a gemma model. get_vertex_ai_model_route must return PARTNER_MODELS, never
GEMMA, MODEL_GARDEN, or NON_GEMINI.
"""
from litellm.llms.vertex_ai.common_utils import (
VertexAIModelRoute,
get_vertex_ai_model_route,
)
route = get_vertex_ai_model_route("google/gemma-4-26b-a4b-it-maas")
assert route == VertexAIModelRoute.PARTNER_MODELS
def test_vertex_ai_google_gemini_not_detected_as_gemma_maas():
"""
Negative: adding the "google/gemma-" prefix must not widen detection to
other google/* models like google/gemini-* (which should keep flowing
through the gemini route, not partner_models).
"""
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
VertexAIPartnerModels,
)
assert not VertexAIPartnerModels.is_vertex_partner_model("google/gemini-1.5-pro")
assert not VertexAIPartnerModels.should_use_openai_handler("google/gemini-1.5-pro")
def test_build_vertex_schema_empty_properties():
"""
Test _build_vertex_schema handles empty properties objects correctly.
This test verifies the fix for the issue where Gemini rejects schemas
This test verifies the fix for the issue where Gemini rejects schemas
with empty properties objects like {"properties": {}, "type": "object"}.
Error from Gemini: "GenerateContentRequest.generation_config.response_schema
.properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties:
.properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties:
should be non-empty for OBJECT type"
The fix removes empty properties objects and their associated type/required fields.
"""
from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
@ -1281,20 +1357,20 @@ def test_build_vertex_schema_empty_properties():
"type": "object",
"additionalProperties": False,
"description": "Go back",
"required": []
"required": [],
}
},
"required": ["go_back"],
"type": "object",
"additionalProperties": False
"additionalProperties": False,
}
]
},
"type": "array"
"type": "array",
}
},
"type": "object",
"additionalProperties": False
"additionalProperties": False,
}
# Apply the transformation
@ -1302,24 +1378,36 @@ def test_build_vertex_schema_empty_properties():
# Verify the transformation removed empty properties
# Navigate to the go_back schema
go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"]["go_back"]
go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"][
"go_back"
]
# Verify empty properties was removed
assert "properties" not in go_back_schema, "Empty properties should be removed"
# Verify type is kept as object (Gemini requires type: object even without properties)
assert go_back_schema.get("type") == "object", "Type should be kept as object when properties is empty"
assert (
go_back_schema.get("type") == "object"
), "Type should be kept as object when properties is empty"
# Verify required was also removed
assert "required" not in go_back_schema, "Required should be removed when properties is empty"
assert (
"required" not in go_back_schema
), "Required should be removed when properties is empty"
# Verify description is preserved
assert go_back_schema.get("description") == "Go back", "Description should be preserved"
assert (
go_back_schema.get("description") == "Go back"
), "Description should be preserved"
# Verify parent schema still has proper structure
parent_schema = result["properties"]["action"]["items"]["anyOf"][0]
assert parent_schema["type"] == "object", "Parent schema should still have object type"
assert "go_back" in parent_schema["properties"], "go_back should still be in parent properties"
assert (
parent_schema["type"] == "object"
), "Parent schema should still have object type"
assert (
"go_back" in parent_schema["properties"]
), "go_back should still be in parent properties"
def test_add_object_type_schema_with_no_properties_and_no_type():
@ -1330,9 +1418,7 @@ def test_add_object_type_schema_with_no_properties_and_no_type():
from litellm.llms.vertex_ai.common_utils import add_object_type
# Input: Schema with no properties and no type (the problematic case)
input_schema = {
"$schema": "https://json-schema.org/draft/2020-12/schema"
}
input_schema = {"$schema": "https://json-schema.org/draft/2020-12/schema"}
# Apply the transformation
add_object_type(input_schema)
@ -1351,10 +1437,7 @@ def test_add_object_type_does_not_override_existing_type():
from litellm.llms.vertex_ai.common_utils import add_object_type
# Input: Schema with existing type
input_schema = {
"type": "string",
"description": "A string field"
}
input_schema = {"type": "string", "description": "A string field"}
# Apply the transformation
add_object_type(input_schema)
@ -1370,12 +1453,7 @@ def test_add_object_type_does_not_add_type_when_anyof_present():
from litellm.llms.vertex_ai.common_utils import add_object_type
# Input: Schema with anyOf but no type
input_schema = {
"anyOf": [
{"type": "string"},
{"type": "null"}
]
}
input_schema = {"anyOf": [{"type": "string"}, {"type": "null"}]}
# Apply the transformation
add_object_type(input_schema)

0
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/__init__.py Normal file
View file

191
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py Normal file
View file

@ -0,0 +1,191 @@
"""
Tests for Vertex AI Gemma MaaS models that route through the partner-models
OpenAI-compatible path (https://aiplatform.googleapis.com/.../endpoints/openapi).
These tests verify that:
1. The correct global URL is constructed (https://aiplatform.googleapis.com)
2. get_vertex_region resolves to "global" when model_cost says so
3. acompletion() goes through the OpenAI-compatible handler and hits
/endpoints/openapi/chat/completions
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../../../../..")
) # Adds the parent directory to the system path
import litellm
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VertexPartnerProvider
@pytest.fixture(autouse=True)
def clean_vertex_env():
"""Clear Google/Vertex AI environment variables before each test to prevent test isolation issues."""
saved_env = {}
env_vars_to_clear = [
"GOOGLE_APPLICATION_CREDENTIALS",
"GOOGLE_CLOUD_PROJECT",
"VERTEXAI_PROJECT",
"VERTEX_PROJECT",
"VERTEX_LOCATION",
"VERTEX_AI_PROJECT",
]
for var in env_vars_to_clear:
if var in os.environ:
saved_env[var] = os.environ[var]
del os.environ[var]
yield
for var, value in saved_env.items():
os.environ[var] = value
class TestVertexBaseGetVertexRegionGemma:
"""Test the get_vertex_region method for Gemma MaaS via model_cost lookup."""
def test_global_model_no_user_region_returns_global(self):
vertex_base = VertexBase()
with patch.dict(
litellm.model_cost,
{
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
"supported_regions": ["global"]
}
},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region=None,
model="google/gemma-4-26b-a4b-it-maas",
)
assert result == "global"
def test_global_model_with_unsupported_user_region_overrides(self):
vertex_base = VertexBase()
with patch.dict(
litellm.model_cost,
{
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
"supported_regions": ["global"]
}
},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region="us-central1",
model="google/gemma-4-26b-a4b-it-maas",
)
assert result == "global"
class TestCreateVertexURLGemma:
"""Test that create_vertex_url produces the expected OpenAI-compatible URL."""
def test_global_location_url_format(self):
url = VertexBase.create_vertex_url(
vertex_location="global",
vertex_project="test-project",
partner=VertexPartnerProvider.llama,
stream=False,
model="google/gemma-4-26b-a4b-it-maas",
)
assert url.startswith("https://aiplatform.googleapis.com")
assert "global-aiplatform.googleapis.com" not in url
assert "/locations/global/" in url
assert url.endswith("/endpoints/openapi/chat/completions")
def test_regional_location_url_format(self):
url = VertexBase.create_vertex_url(
vertex_location="us-central1",
vertex_project="test-project",
partner=VertexPartnerProvider.llama,
stream=False,
model="google/gemma-4-26b-a4b-it-maas",
)
assert url.startswith("https://us-central1-aiplatform.googleapis.com")
assert "/locations/us-central1/" in url
assert url.endswith("/endpoints/openapi/chat/completions")
@pytest.mark.asyncio
async def test_vertex_ai_gemma_global_endpoint_url():
"""
End-to-end: acompletion on vertex_ai/google/gemma-4-26b-a4b-it-maas should
POST to the global endpoints/openapi/chat/completions URL.
"""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.json.return_value = {
"id": "chatcmpl-gemma-test",
"object": "chat.completion",
"created": 1234567890,
"model": "google/gemma-4-26b-a4b-it-maas",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "Hello! How can I help you today?",
},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
}
mock_vertexai = MagicMock()
mock_vertexai.preview = MagicMock()
with (
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler"
) as mock_http_handler,
patch(
"litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token",
return_value=("fake-token", "test-project"),
),
patch.dict(
"sys.modules",
{"vertexai": mock_vertexai, "vertexai.preview": mock_vertexai.preview},
),
patch.dict(
litellm.model_cost,
{
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
"supported_regions": ["global"]
}
},
clear=False,
),
):
mock_http_handler.return_value.post = AsyncMock(return_value=mock_response)
response = await litellm.acompletion(
model="vertex_ai/google/gemma-4-26b-a4b-it-maas",
messages=[{"role": "user", "content": "Hello"}],
vertex_ai_project="test-project",
)
mock_http_handler.return_value.post.assert_called_once()
call_args = mock_http_handler.return_value.post.call_args
called_url = call_args.kwargs["url"]
assert called_url.startswith("https://aiplatform.googleapis.com")
assert "global-aiplatform.googleapis.com" not in called_url
assert "/locations/global/" in called_url
assert "/endpoints/openapi/chat/completions" in called_url
assert response.model == "google/gemma-4-26b-a4b-it-maas"