Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 54 additions & 0 deletions litellm/llms/vertex_ai/common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,57 @@ def _fix_enum_empty_strings(schema, depth=0):
_fix_enum_empty_strings(items, depth=depth + 1)


def _fix_enum_types(schema, depth=0):
"""Remove `enum` fields when the schema type is not string.

Gemini / Vertex APIs only allow enums for string-typed fields. When an enum
is present on a non-string typed property (or when `anyOf` types do not
include a string type), remove the enum to avoid provider validation errors.
"""
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError(
f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema."
)

if not isinstance(schema, dict):
return

# If enum exists but type is not string (and anyOf doesn't include string), drop enum
if "enum" in schema and isinstance(schema["enum"], list):
schema_type = schema.get("type")
keep_enum = False
if isinstance(schema_type, str) and schema_type.lower() == "string":
keep_enum = True
else:
anyof = schema.get("anyOf")
if isinstance(anyof, list):
for item in anyof:
if isinstance(item, dict):
item_type = item.get("type")
if isinstance(item_type, str) and item_type.lower() == "string":
keep_enum = True
break

if not keep_enum:
schema.pop("enum", None)

# Recurse into nested structures
properties = schema.get("properties", None)
if properties is not None:
for _, value in properties.items():
_fix_enum_types(value, depth=depth + 1)

items = schema.get("items", None)
if items is not None:
_fix_enum_types(items, depth=depth + 1)

anyof = schema.get("anyOf", None)
if anyof is not None and isinstance(anyof, list):
for item in anyof:
if isinstance(item, dict):
_fix_enum_types(item, depth=depth + 1)


def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
"""
This is a modified version of https://github.com/google-gemini/generative-ai-python/blob/8f77cc6ac99937cd3a81299ecf79608b91b06bbb/google/generativeai/types/content_types.py#L419
Expand Down Expand Up @@ -307,6 +358,9 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
# Handle empty strings in enum values - Gemini doesn't accept empty strings in enums
_fix_enum_empty_strings(parameters)

# Remove enums for non-string typed fields (Gemini requires enum only on strings)
_fix_enum_types(parameters)

# Handle empty items objects
process_items(parameters)
add_object_type(parameters)
Expand Down
118 changes: 118 additions & 0 deletions tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -803,6 +803,124 @@ def test_fix_enum_empty_strings():
assert input_schema["properties"]["user_agent_type"]["description"] == "Device type for user agent"


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.

Relevant issue: Gemini only allows enums for string-typed fields
"""
from litellm.llms.vertex_ai.common_utils import _fix_enum_types

# Input: Schema with enum on non-string type (the problematic case)
input_schema = {
"type": "object",
"properties": {
"truncateMode": {
"enum": ["auto", "none", "start", "end"],
"type": "string", # This should keep the enum
"description": "How to truncate content"
},
"maxLength": {
"enum": [100, 200, 500], # This should be removed
"type": "integer",
"description": "Maximum length"
},
"enabled": {
"enum": [True, False], # This should be removed
"type": "boolean",
"description": "Whether feature is enabled"
},
"nested": {
"type": "object",
"properties": {
"innerEnum": {
"enum": ["a", "b", "c"], # This should be kept
"type": "string"
},
"innerNonStringEnum": {
"enum": [1, 2, 3], # This should be removed
"type": "integer"
}
}
},
"anyOfField": {
"anyOf": [
{"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
expected_output = {
"type": "object",
"properties": {
"truncateMode": {
"enum": ["auto", "none", "start", "end"], # Kept - string type
"type": "string",
"description": "How to truncate content"
},
"maxLength": { # enum removed
"type": "integer",
"description": "Maximum length"
},
"enabled": { # enum removed
"type": "boolean",
"description": "Whether feature is enabled"
},
"nested": {
"type": "object",
"properties": {
"innerEnum": {
"enum": ["a", "b", "c"], # Kept - string type
"type": "string"
},
"innerNonStringEnum": { # enum removed
"type": "integer"
}
}
},
"anyOfField": {
"anyOf": [
{"type": "string", "enum": ["option1", "option2"]}, # Kept - has string type
{"type": "integer"} # enum removed
]
}
}
}

# Apply the transformation
_fix_enum_types(input_schema)

# Verify the transformation
assert input_schema == expected_output

# 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 "enum" in input_schema["properties"]["nested"]["properties"]["innerEnum"]
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"]

# 3. anyOf with string type keeps enum, non-string removes it
assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0]
assert "enum" not in input_schema["properties"]["anyOfField"]["anyOf"][1]

# 4. Other properties preserved
assert input_schema["properties"]["maxLength"]["type"] == "integer"
assert input_schema["properties"]["enabled"]["type"] == "boolean"


def test_get_token_url():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
Expand Down
Loading