diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index 6cd3aceb079..cd27b4c362a 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -480,6 +480,62 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
else:
return None
+ @staticmethod
+ def _resolve_search_tool_conflict(
+ gtool_func_declarations: list,
+ googleSearch: Optional[dict],
+ googleSearchRetrieval: Optional[dict],
+ enterpriseWebSearch: Optional[dict],
+ urlContext: Optional[dict],
+ optional_params: dict,
+ ) -> tuple:
+ """
+ Resolve Vertex AI constraint: multiple Tool objects in a request must
+ ALL be search tools. When function declarations are mixed with search
+ tools, drop search tools to avoid 400 error.
+
+ Skip when include_server_side_tool_invocations is enabled (Gemini 3+
+ supports tool combination natively).
+
+ Note: code_execution, computerUse, and googleMaps are NOT search tools
+ and CAN coexist with function declarations, so they are preserved.
+
+ Ref: https://github.com/BerriAI/litellm/issues/23337
+
+ Returns:
+ tuple of (googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext)
+ """
+ has_search_tools = any(
+ v is not None
+ for v in [
+ googleSearch,
+ googleSearchRetrieval,
+ enterpriseWebSearch,
+ urlContext,
+ ]
+ )
+ server_side_tool_invocations = optional_params.get(
+ "include_server_side_tool_invocations", False
+ )
+ if (
+ gtool_func_declarations
+ and has_search_tools
+ and not server_side_tool_invocations
+ ):
+ verbose_logger.warning(
+ "Vertex AI does not support mixing function declarations with "
+ "search tools (googleSearch, enterpriseWebSearch, urlContext, "
+ "googleSearchRetrieval) in the same request. Dropping search "
+ "tools and keeping function declarations. To use search tools, "
+ "send a request without function calling tools."
+ )
+ googleSearch = None
+ googleSearchRetrieval = None
+ enterpriseWebSearch = None
+ urlContext = None
+
+ return googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext
+
def _map_function( # noqa: PLR0915
self, value: List[dict], optional_params: dict
) -> List[Tools]:
@@ -512,9 +568,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
value = _remove_strict_from_schema(value)
for tool in value:
- openai_function_object: Optional[
- ChatCompletionToolParamFunctionChunk
- ] = None
+ openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = (
+ None
+ )
if "function" in tool: # tools list
_openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore
**tool["function"]
@@ -633,43 +689,19 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
# per Vertex AI API spec: "A Tool object should contain exactly one type of Tool"
_tools_list: List[Tools] = []
- # Vertex AI constraint: multiple Tool objects in a request must ALL be
- # search tools. Mixing function declarations with search tools in the
- # same request causes a 400 error:
- # "Multiple tools are supported only when they are all search tools."
- # When both are present (e.g. deployment config has search tools and
- # user request adds function calling tools via MCP), drop search tools
- # and keep function declarations.
- # Ref: https://github.com/BerriAI/litellm/issues/23337
- has_search_tools = any(
- v is not None
- for v in [
- googleSearch,
- googleSearchRetrieval,
- enterpriseWebSearch,
- urlContext,
- ]
+ (
+ googleSearch,
+ googleSearchRetrieval,
+ enterpriseWebSearch,
+ urlContext,
+ ) = self._resolve_search_tool_conflict(
+ gtool_func_declarations=gtool_func_declarations,
+ googleSearch=googleSearch,
+ googleSearchRetrieval=googleSearchRetrieval,
+ enterpriseWebSearch=enterpriseWebSearch,
+ urlContext=urlContext,
+ optional_params=optional_params,
)
- # Skip this check when include_server_side_tool_invocations is enabled
- # (Gemini 3+ supports tool combination natively via PR #24073).
- server_side_tool_invocations = optional_params.get(
- "include_server_side_tool_invocations", False
- )
- if gtool_func_declarations and has_search_tools and not server_side_tool_invocations:
- verbose_logger.warning(
- "Vertex AI does not support mixing function declarations with "
- "search tools (googleSearch, enterpriseWebSearch, urlContext, "
- "googleSearchRetrieval) in the same request. Dropping search "
- "tools and keeping function declarations. To use search tools, "
- "send a request without function calling tools."
- )
- googleSearch = None
- googleSearchRetrieval = None
- enterpriseWebSearch = None
- urlContext = None
- # Note: code_execution, computerUse, and googleMaps are NOT search
- # tools and CAN coexist with function declarations in separate Tool
- # objects, so they are intentionally preserved here.
# Function declarations can be grouped together in one Tool
if gtool_func_declarations:
@@ -684,15 +716,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_tools_list.append(search_tool)
if googleSearchRetrieval is not None:
retrieval_tool = Tools()
- retrieval_tool[
- VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
- ] = googleSearchRetrieval
+ retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = (
+ googleSearchRetrieval
+ )
_tools_list.append(retrieval_tool)
if enterpriseWebSearch is not None:
enterprise_tool = Tools()
- enterprise_tool[
- VertexToolName.ENTERPRISE_WEB_SEARCH.value
- ] = enterpriseWebSearch
+ enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = (
+ enterpriseWebSearch
+ )
_tools_list.append(enterprise_tool)
if code_execution is not None:
code_tool = Tools()
@@ -1139,16 +1171,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
param_description="thinking_budget",
)
if VertexGeminiConfig._is_gemini_3_or_newer(model):
- optional_params[
- "thinkingConfig"
- ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
- effort_value, model
+ optional_params["thinkingConfig"] = (
+ VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
+ effort_value, model
+ )
)
else:
- optional_params[
- "thinkingConfig"
- ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(
- effort_value, model
+ optional_params["thinkingConfig"] = (
+ VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(
+ effort_value, model
+ )
)
elif param == "thinking":
# Validate no conflict with thinking_level
@@ -1157,11 +1189,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
param_name="thinking",
param_description="thinking_budget",
)
- optional_params[
- "thinkingConfig"
- ] = VertexGeminiConfig._map_thinking_param(
- cast(AnthropicThinkingParam, value),
- model=model,
+ optional_params["thinkingConfig"] = (
+ VertexGeminiConfig._map_thinking_param(
+ cast(AnthropicThinkingParam, value),
+ model=model,
+ )
)
elif param == "modalities" and isinstance(value, list):
response_modalities = self.map_response_modalities(value)
@@ -1585,10 +1617,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_tool_response_chunk["provider_specific_fields"] = { # type: ignore
"thought_signature": thought_signature
}
- _tool_response_chunk[
- "id"
- ] = _encode_tool_call_id_with_signature(
- _tool_response_chunk["id"] or "", thought_signature
+ _tool_response_chunk["id"] = (
+ _encode_tool_call_id_with_signature(
+ _tool_response_chunk["id"] or "", thought_signature
+ )
)
_tools.append(_tool_response_chunk)
cumulative_tool_call_idx += 1
@@ -2435,28 +2467,28 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
## ADD METADATA TO RESPONSE ##
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata)
- model_response._hidden_params[
- "vertex_ai_grounding_metadata"
- ] = grounding_metadata
+ model_response._hidden_params["vertex_ai_grounding_metadata"] = (
+ grounding_metadata
+ )
setattr(
model_response, "vertex_ai_url_context_metadata", url_context_metadata
)
- model_response._hidden_params[
- "vertex_ai_url_context_metadata"
- ] = url_context_metadata
+ model_response._hidden_params["vertex_ai_url_context_metadata"] = (
+ url_context_metadata
+ )
setattr(model_response, "vertex_ai_safety_results", safety_ratings)
- model_response._hidden_params[
- "vertex_ai_safety_results"
- ] = safety_ratings # older approach - maintaining to prevent regressions
+ model_response._hidden_params["vertex_ai_safety_results"] = (
+ safety_ratings # older approach - maintaining to prevent regressions
+ )
## ADD CITATION METADATA ##
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata)
- model_response._hidden_params[
- "vertex_ai_citation_metadata"
- ] = citation_metadata # older approach - maintaining to prevent regressions
+ model_response._hidden_params["vertex_ai_citation_metadata"] = (
+ citation_metadata # older approach - maintaining to prevent regressions
+ )
## ADD TRAFFIC TYPE ##
traffic_type = completion_response.get("usageMetadata", {}).get(
@@ -3164,7 +3196,12 @@ class ModelResponseIterator:
setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore
- return grounding_metadata, url_context_metadata, safety_ratings, citation_metadata
+ return (
+ grounding_metadata,
+ url_context_metadata,
+ safety_ratings,
+ citation_metadata,
+ )
def _apply_stream_usage_metadata(
self,
@@ -3189,9 +3226,9 @@ class ModelResponseIterator:
traffic_type = processed_chunk.get("usageMetadata", {}).get("trafficType")
if traffic_type:
- model_response._hidden_params.setdefault(
- "provider_specific_fields", {}
- )["traffic_type"] = traffic_type
+ model_response._hidden_params.setdefault("provider_specific_fields", {})[
+ "traffic_type"
+ ] = traffic_type
service_tier = self.response_headers.get("x-gemini-service-tier")
if service_tier:
diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
index 2e719b212f7..a0979664943 100644
--- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
@@ -237,7 +237,9 @@ def test_vertex_ai_response_json_schema_preserves_refs_for_gemini_2():
# $defs and $ref should be preserved (not unpacked)
assert "response_json_schema" in transformed_request
result_schema = transformed_request["response_json_schema"]
- assert "$defs" in result_schema, "responseJsonSchema should preserve $defs for Gemini 2.0+"
+ assert (
+ "$defs" in result_schema
+ ), "responseJsonSchema should preserve $defs for Gemini 2.0+"
def test_vertex_ai_get_json_schema_preserves_refs_for_nested_pydantic():
@@ -317,14 +319,22 @@ def test_vertex_ai_response_json_schema_for_gemini_2():
# Types should be lowercase (standard JSON Schema format)
assert transformed_request["response_json_schema"]["type"] == "object"
- assert transformed_request["response_json_schema"]["properties"]["name"]["type"] == "string"
- assert transformed_request["response_json_schema"]["properties"]["age"]["type"] == "integer"
+ assert (
+ transformed_request["response_json_schema"]["properties"]["name"]["type"]
+ == "string"
+ )
+ assert (
+ transformed_request["response_json_schema"]["properties"]["age"]["type"]
+ == "integer"
+ )
# Should NOT have propertyOrdering (not needed for responseJsonSchema)
assert "propertyOrdering" not in transformed_request["response_json_schema"]
# additionalProperties should be preserved (supported by responseJsonSchema)
- assert transformed_request["response_json_schema"].get("additionalProperties") == False
+ assert (
+ transformed_request["response_json_schema"].get("additionalProperties") == False
+ )
def test_vertex_ai_response_schema_for_old_models():
@@ -581,7 +591,7 @@ def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content(
"args": {"timezone": "America/New_York"},
},
"thoughtSignature": "EsEDCr4DAdHtim...", # Just a token, not reasoning
- }
+ },
]
},
"finishReason": "STOP",
@@ -600,12 +610,18 @@ def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content(
streaming_chunk = iterator.chunk_parser(chunk)
# Verify reasoning_content comes from the thought: true part
- assert streaming_chunk.choices[0].delta.reasoning_content == "Let me think about how to get the time..."
+ assert (
+ streaming_chunk.choices[0].delta.reasoning_content
+ == "Let me think about how to get the time..."
+ )
# Verify tool calls are also present
assert streaming_chunk.choices[0].delta.tool_calls is not None
assert len(streaming_chunk.choices[0].delta.tool_calls) == 1
- assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time"
+ assert (
+ streaming_chunk.choices[0].delta.tool_calls[0].function.name
+ == "get_current_time"
+ )
def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content():
@@ -653,12 +669,15 @@ def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content():
streaming_chunk = iterator.chunk_parser(chunk)
# reasoning_content should be None - thoughtSignature alone does NOT mean reasoning
- assert getattr(streaming_chunk.choices[0].delta, 'reasoning_content', None) is None
+ assert getattr(streaming_chunk.choices[0].delta, "reasoning_content", None) is None
# Tool calls should still work
assert streaming_chunk.choices[0].delta.tool_calls is not None
assert len(streaming_chunk.choices[0].delta.tool_calls) == 1
- assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time"
+ assert (
+ streaming_chunk.choices[0].delta.tool_calls[0].function.name
+ == "get_current_time"
+ )
def test_check_finish_reason():
@@ -711,7 +730,10 @@ def test_vertex_ai_usage_metadata_response_token_count():
"promptTokenCount": 66,
"responseTokenCount": 74,
"totalTokenCount": 131,
- "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 57}, {"modality": "IMAGE", "tokenCount": 9}],
+ "promptTokensDetails": [
+ {"modality": "TEXT", "tokenCount": 57},
+ {"modality": "IMAGE", "tokenCount": 9},
+ ],
"responseTokensDetails": [{"modality": "TEXT", "tokenCount": 74}],
}
usage_metadata = UsageMetadata(**usage_metadata)
@@ -741,9 +763,9 @@ def test_vertex_ai_usage_metadata_with_image_tokens():
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 14}],
"candidatesTokensDetails": [
{"modality": "IMAGE", "tokenCount": 1120},
- {"modality": "TEXT", "tokenCount": 322} # 1442 - 1120 = 322
+ {"modality": "TEXT", "tokenCount": 322}, # 1442 - 1120 = 322
],
- "thoughtsTokenCount": 158
+ "thoughtsTokenCount": 158,
}
usage_metadata = UsageMetadata(**usage_metadata)
result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata})
@@ -785,7 +807,7 @@ def test_vertex_ai_usage_metadata_with_image_tokens_auto_calculated_text():
{"modality": "IMAGE", "tokenCount": 1120}
# TEXT modality omitted - should be auto-calculated
],
- "thoughtsTokenCount": 158
+ "thoughtsTokenCount": 158,
}
usage_metadata = UsageMetadata(**usage_metadata)
result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata})
@@ -809,13 +831,13 @@ def test_vertex_ai_usage_metadata_with_image_tokens_auto_calculated_text():
def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt():
"""Test promptTokensDetails with IMAGE modality for multimodal inputs
-
+
This test verifies the fix for issue #18182 where image_tokens were missing
from prompt_tokens_details when calling Gemini models with image inputs.
-
+
Example scenario: User sends a text prompt + image, and Gemini generates an image response.
The promptTokensDetails should include both TEXT and IMAGE token counts.
-
+
In this test case, candidatesTokenCount is INCLUSIVE of thoughtsTokenCount because:
promptTokenCount (533) + candidatesTokenCount (1337) = totalTokenCount (1870)
"""
@@ -826,31 +848,29 @@ def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt():
"totalTokenCount": 1870,
"promptTokensDetails": [
{"modality": "IMAGE", "tokenCount": 527},
- {"modality": "TEXT", "tokenCount": 6}
+ {"modality": "TEXT", "tokenCount": 6},
],
- "candidatesTokensDetails": [
- {"modality": "IMAGE", "tokenCount": 1120}
- ],
- "thoughtsTokenCount": 217
+ "candidatesTokensDetails": [{"modality": "IMAGE", "tokenCount": 1120}],
+ "thoughtsTokenCount": 217,
}
usage_metadata = UsageMetadata(**usage_metadata)
result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata})
print("result", result)
-
+
# Verify basic token counts
assert result.prompt_tokens == 533
# candidatesTokenCount is INCLUSIVE, so completion_tokens = candidatesTokenCount
assert result.completion_tokens == 1337
assert result.total_tokens == 1870
-
+
# Verify prompt_tokens_details includes both text and image tokens
assert result.prompt_tokens_details.text_tokens == 6
assert result.prompt_tokens_details.image_tokens == 527
-
+
# Verify completion_tokens_details
assert result.completion_tokens_details.image_tokens == 1120
assert result.completion_tokens_details.reasoning_tokens == 217
-
+
# Verify the math: prompt_tokens = text + image
# 533 = 6 (text) + 527 (image)
assert (
@@ -916,13 +936,17 @@ def test_vertex_ai_map_thinking_param_with_budget_tokens_0():
def test_vertex_ai_map_tools():
v = VertexGeminiConfig()
optional_params = {}
- tools = v._map_function(value=[{"code_execution": {}}], optional_params=optional_params)
+ tools = v._map_function(
+ value=[{"code_execution": {}}], optional_params=optional_params
+ )
assert len(tools) == 1
assert tools[0]["code_execution"] == {}
print(tools)
new_optional_params = {}
- new_tools = v._map_function(value=[{"codeExecution": {}}], optional_params=new_optional_params)
+ new_tools = v._map_function(
+ value=[{"codeExecution": {}}], optional_params=new_optional_params
+ )
assert len(new_tools) == 1
print("new_tools", new_tools)
assert new_tools[0]["code_execution"] == {}
@@ -1088,7 +1112,13 @@ def test_vertex_ai_streaming_usage_web_search_calculation():
{
"content": {"parts": [{"text": "Hello"}]},
"groundingMetadata": [
- {"webSearchQueries": ["", "What is the capital of France?", "Capital of France"]}
+ {
+ "webSearchQueries": [
+ "",
+ "What is the capital of France?",
+ "Capital of France",
+ ]
+ }
],
}
],
@@ -1432,7 +1462,7 @@ def test_vertex_ai_process_candidates_with_grounding_metadata():
def test_vertex_ai_tool_call_id_format():
"""
Test that tool call IDs have the correct format and length.
-
+
The ID should be in format 'call_' + 28 hex characters (total 33 characters).
This test verifies the fix for keeping the code line under 40 characters.
"""
@@ -1449,12 +1479,7 @@ def test_vertex_ai_tool_call_id_format():
"args": {"location": "San Francisco", "unit": "celsius"},
}
),
- HttpxPartType(
- functionCall={
- "name": "get_time",
- "args": {"timezone": "PST"}
- }
- ),
+ HttpxPartType(functionCall={"name": "get_time", "args": {"timezone": "PST"}}),
]
function, tools, updated_idx = VertexGeminiConfig._transform_parts(
@@ -1469,19 +1494,27 @@ def test_vertex_ai_tool_call_id_format():
# Test ID format for both tool calls
for tool in tools:
tool_id = tool["id"]
-
+
# Should start with 'call_'
- assert tool_id.startswith("call_"), f"ID should start with 'call_', got: {tool_id}"
-
+ assert tool_id.startswith(
+ "call_"
+ ), f"ID should start with 'call_', got: {tool_id}"
+
# Should have exactly 33 total characters (call_ + 28 hex chars)
- assert len(tool_id) == 33, f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}"
-
+ assert (
+ len(tool_id) == 33
+ ), f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}"
+
# The part after 'call_' should be 28 hex characters
hex_part = tool_id[5:] # Remove 'call_' prefix
- assert len(hex_part) == 28, f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}"
-
+ assert (
+ len(hex_part) == 28
+ ), f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}"
+
# Should only contain valid hex characters
- assert re.match(r'^[0-9a-f]{28}$', hex_part), f"Should contain only lowercase hex chars, got: {hex_part}"
+ assert re.match(
+ r"^[0-9a-f]{28}$", hex_part
+ ), f"Should contain only lowercase hex chars, got: {hex_part}"
# Verify IDs are unique
assert tools[0]["id"] != tools[1]["id"], "Tool call IDs should be unique"
@@ -1496,15 +1529,17 @@ def test_vertex_ai_tool_call_id_format():
)
if test_tools:
ids_generated.add(test_tools[0]["id"])
-
+
# All generated IDs should be unique
- assert len(ids_generated) == 10, f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs"
+ assert (
+ len(ids_generated) == 10
+ ), f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs"
def test_vertex_ai_code_line_length():
"""
Test that the specific code line generating tool call IDs is within character limit.
-
+
This is a meta-test to ensure the code change meets the 40-character requirement.
"""
import inspect
@@ -1514,45 +1549,49 @@ def test_vertex_ai_code_line_length():
)
# Get the source code of the _transform_parts method
- source_lines = inspect.getsource(VertexGeminiConfig._transform_parts).split('\n')
-
+ source_lines = inspect.getsource(VertexGeminiConfig._transform_parts).split("\n")
+
# Find the line that generates the ID
id_line = None
for line in source_lines:
- if '"id": f"call_' in line and 'uuid.uuid4().hex[:28]' in line:
+ if '"id": f"call_' in line and "uuid.uuid4().hex[:28]" in line:
id_line = line.strip() # Remove indentation for length check
break
-
+
assert id_line is not None, "Could not find the ID generation line in source code"
-
+
# Check that the line is 40 characters or less (excluding indentation)
line_length = len(id_line)
- assert line_length <= 40, f"ID generation line is {line_length} characters, should be ≤40: {id_line}"
-
+ assert (
+ line_length <= 40
+ ), f"ID generation line is {line_length} characters, should be ≤40: {id_line}"
+
# Verify it contains the expected UUID format
- assert 'uuid.uuid4().hex[:28]' in id_line, f"Line should contain shortened UUID format: {id_line}"
+ assert (
+ "uuid.uuid4().hex[:28]" in id_line
+ ), f"Line should contain shortened UUID format: {id_line}"
def test_vertex_ai_map_google_maps_tool_simple():
"""
Test googleMaps tool transformation without location data.
-
+
Input:
value=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}]
optional_params={}
-
+
Expected Output:
tools=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}]
optional_params={} (unchanged)
"""
v = VertexGeminiConfig()
optional_params = {}
-
+
tools = v._map_function(
value=[{"googleMaps": {"enableWidget": "ENABLE_WIDGET"}}],
- optional_params=optional_params
+ optional_params=optional_params,
)
-
+
assert len(tools) == 1
assert "googleMaps" in tools[0]
assert tools[0]["googleMaps"]["enableWidget"] == "ENABLE_WIDGET"
@@ -1563,7 +1602,7 @@ def test_vertex_ai_map_google_maps_tool_with_location():
"""
Test googleMaps tool transformation with location data.
Verifies latitude/longitude/languageCode are extracted to toolConfig.retrievalConfig.
-
+
Input:
value=[{
"googleMaps": {
@@ -1574,7 +1613,7 @@ def test_vertex_ai_map_google_maps_tool_with_location():
}
}]
optional_params={}
-
+
Expected Output:
tools=[{
"googleMaps": {"enableWidget": "ENABLE_WIDGET"}
@@ -1593,40 +1632,43 @@ def test_vertex_ai_map_google_maps_tool_with_location():
"""
v = VertexGeminiConfig()
optional_params = {}
-
+
tools = v._map_function(
- value=[{
- "googleMaps": {
- "enableWidget": "ENABLE_WIDGET",
- "latitude": 37.7749,
- "longitude": -122.4194,
- "languageCode": "en_US"
+ value=[
+ {
+ "googleMaps": {
+ "enableWidget": "ENABLE_WIDGET",
+ "latitude": 37.7749,
+ "longitude": -122.4194,
+ "languageCode": "en_US",
+ }
}
- }],
- optional_params=optional_params
+ ],
+ optional_params=optional_params,
)
-
+
assert len(tools) == 1
assert "googleMaps" in tools[0]
-
+
google_maps_tool = tools[0]["googleMaps"]
assert google_maps_tool["enableWidget"] == "ENABLE_WIDGET"
assert "latitude" not in google_maps_tool
assert "longitude" not in google_maps_tool
assert "languageCode" not in google_maps_tool
-
+
assert "toolConfig" in optional_params
assert "retrievalConfig" in optional_params["toolConfig"]
-
+
retrieval_config = optional_params["toolConfig"]["retrievalConfig"]
assert retrieval_config["latLng"]["latitude"] == 37.7749
assert retrieval_config["latLng"]["longitude"] == -122.4194
assert retrieval_config["languageCode"] == "en_US"
+
def test_vertex_ai_penalty_parameters_validation():
"""
Test that penalty parameters are properly validated for different Gemini models.
-
+
This test ensures that:
1. Models that don't support penalty parameters (like preview models) filter them out
2. Models that support penalty parameters include them in the request
@@ -1641,14 +1683,19 @@ def test_vertex_ai_penalty_parameters_validation():
for model, should_support in test_cases:
# Test _supports_penalty_parameters method
- assert v._supports_penalty_parameters(model) == should_support, \
- f"Model {model} penalty support should be {should_support}"
+ assert (
+ v._supports_penalty_parameters(model) == should_support
+ ), f"Model {model} penalty support should be {should_support}"
# Test get_supported_openai_params method
supported_params = v.get_supported_openai_params(model)
- has_penalty_params = "frequency_penalty" in supported_params and "presence_penalty" in supported_params
- assert has_penalty_params == should_support, \
- f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list"
+ has_penalty_params = (
+ "frequency_penalty" in supported_params
+ and "presence_penalty" in supported_params
+ )
+ assert (
+ has_penalty_params == should_support
+ ), f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list"
# Test parameter mapping for unsupported model
model = "gemini-2.5-pro-preview-06-05"
@@ -1656,7 +1703,7 @@ def test_vertex_ai_penalty_parameters_validation():
"temperature": 0.7,
"frequency_penalty": 0.5,
"presence_penalty": 0.3,
- "max_tokens": 100
+ "max_tokens": 100,
}
optional_params = {}
@@ -1664,12 +1711,16 @@ def test_vertex_ai_penalty_parameters_validation():
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
- drop_params=False
+ drop_params=False,
)
# Penalty parameters should be filtered out for unsupported models
- assert "frequency_penalty" not in result, "frequency_penalty should be filtered out for unsupported model"
- assert "presence_penalty" not in result, "presence_penalty should be filtered out for unsupported model"
+ assert (
+ "frequency_penalty" not in result
+ ), "frequency_penalty should be filtered out for unsupported model"
+ assert (
+ "presence_penalty" not in result
+ ), "presence_penalty should be filtered out for unsupported model"
# Other parameters should still be included
assert "temperature" in result, "temperature should still be included"
@@ -1681,7 +1732,7 @@ def test_vertex_ai_penalty_parameters_validation():
def test_vertex_ai_gemini_3_penalty_parameters_unsupported():
"""
Test that penalty parameters are not supported for Gemini 3 models.
-
+
This test ensures that:
1. Gemini 3 models do not support penalty parameters
2. Penalty parameters are excluded from supported params list for Gemini 3 models
@@ -1698,22 +1749,25 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported():
for model in gemini_3_models:
# Test _supports_penalty_parameters method
- assert v._supports_penalty_parameters(model) == False, \
- f"Gemini 3 model {model} should not support penalty parameters"
+ assert (
+ v._supports_penalty_parameters(model) == False
+ ), f"Gemini 3 model {model} should not support penalty parameters"
# Test get_supported_openai_params method
supported_params = v.get_supported_openai_params(model)
- assert "frequency_penalty" not in supported_params, \
- f"frequency_penalty should not be in supported params for {model}"
- assert "presence_penalty" not in supported_params, \
- f"presence_penalty should not be in supported params for {model}"
+ assert (
+ "frequency_penalty" not in supported_params
+ ), f"frequency_penalty should not be in supported params for {model}"
+ assert (
+ "presence_penalty" not in supported_params
+ ), f"presence_penalty should not be in supported params for {model}"
# Test parameter mapping - penalty params should be filtered out
non_default_params = {
"temperature": 0.7,
"frequency_penalty": 0.5,
"presence_penalty": 0.3,
- "max_tokens": 100
+ "max_tokens": 100,
}
optional_params = {}
@@ -1721,39 +1775,46 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported():
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
- drop_params=False
+ drop_params=False,
)
# Penalty parameters should be filtered out for Gemini 3 models
- assert "frequency_penalty" not in result, \
- f"frequency_penalty should be filtered out for Gemini 3 model {model}"
- assert "presence_penalty" not in result, \
- f"presence_penalty should be filtered out for Gemini 3 model {model}"
+ assert (
+ "frequency_penalty" not in result
+ ), f"frequency_penalty should be filtered out for Gemini 3 model {model}"
+ assert (
+ "presence_penalty" not in result
+ ), f"presence_penalty should be filtered out for Gemini 3 model {model}"
# Other parameters should still be included
- assert "temperature" in result, \
- f"temperature should still be included for Gemini 3 model {model}"
- assert "max_output_tokens" in result, \
- f"max_output_tokens should still be included for Gemini 3 model {model}"
+ assert (
+ "temperature" in result
+ ), f"temperature should still be included for Gemini 3 model {model}"
+ assert (
+ "max_output_tokens" in result
+ ), f"max_output_tokens should still be included for Gemini 3 model {model}"
assert result["temperature"] == 0.7
assert result["max_output_tokens"] == 100
# Test that non-Gemini 3 models still support penalty parameters (if they're not in the unsupported list)
non_gemini_3_model = "gemini-2.5-pro"
- assert v._supports_penalty_parameters(non_gemini_3_model) == True, \
- f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters"
-
+ assert (
+ v._supports_penalty_parameters(non_gemini_3_model) == True
+ ), f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters"
+
supported_params = v.get_supported_openai_params(non_gemini_3_model)
- assert "frequency_penalty" in supported_params, \
- f"frequency_penalty should be in supported params for {non_gemini_3_model}"
- assert "presence_penalty" in supported_params, \
- f"presence_penalty should be in supported params for {non_gemini_3_model}"
+ assert (
+ "frequency_penalty" in supported_params
+ ), f"frequency_penalty should be in supported params for {non_gemini_3_model}"
+ assert (
+ "presence_penalty" in supported_params
+ ), f"presence_penalty should be in supported params for {non_gemini_3_model}"
def test_vertex_ai_annotation_streaming_events():
"""
Test that annotation events are properly emitted during streaming for Vertex AI Gemini.
-
+
This test verifies:
1. Grounding metadata is converted to annotations in streaming chunks
2. Annotations are included in the delta of streaming chunks
@@ -1776,7 +1837,7 @@ def test_vertex_ai_annotation_streaming_events():
"groundingMetadata": {
"webSearchQueries": ["weather San Francisco today"],
"searchEntryPoint": {
- "renderedContent": 'Search results
'
+ "renderedContent": "Search results
"
},
"groundingChunks": [
{
@@ -1817,7 +1878,7 @@ def test_vertex_ai_annotation_streaming_events():
# Verify the chunk was parsed correctly
assert streaming_chunk.choices is not None
assert len(streaming_chunk.choices) == 1
-
+
# Check that annotations are present in the delta
delta = streaming_chunk.choices[0].delta
assert hasattr(delta, "annotations")
@@ -1870,7 +1931,7 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped():
def test_vertex_ai_annotation_conversion():
"""
Test the conversion of Vertex AI grounding metadata to OpenAI annotations.
-
+
This test verifies the _convert_grounding_metadata_to_annotations method
correctly transforms grounding metadata into the expected format.
"""
@@ -1881,9 +1942,7 @@ def test_vertex_ai_annotation_conversion():
# Sample grounding metadata as returned by Vertex AI
grounding_metadata = {
"webSearchQueries": ["weather San Francisco", "current time San Francisco"],
- "searchEntryPoint": {
- "renderedContent": 'Search interface
'
- },
+ "searchEntryPoint": {"renderedContent": "Search interface
"},
"groundingChunks": [
{
"web": {
@@ -1898,7 +1957,7 @@ def test_vertex_ai_annotation_conversion():
"title": "Current time in San Francisco, CA",
"domain": "google.com",
}
- }
+ },
],
"groundingSupports": [
{
@@ -1927,12 +1986,14 @@ def test_vertex_ai_annotation_conversion():
},
"groundingChunkIndices": [1],
"confidenceScores": [0.92],
- }
+ },
],
}
# Convert grounding metadata to annotations
- content_text = "The weather in San Francisco is currently 72°F and the time is 2:30 PM"
+ content_text = (
+ "The weather in San Francisco is currently 72°F and the time is 2:30 PM"
+ )
annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
[grounding_metadata], content_text
)
@@ -1968,7 +2029,7 @@ def test_vertex_ai_annotation_conversion():
def test_vertex_ai_annotation_empty_grounding_metadata():
"""
Test handling of empty or missing grounding metadata.
-
+
This test ensures the annotation conversion handles edge cases gracefully.
"""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@@ -2006,6 +2067,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata():
# ==================== Gemini 3 Pro Preview Tests ====================
+
def test_is_gemini_3_or_newer():
"""Test the _is_gemini_3_or_newer method for version detection"""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@@ -2016,8 +2078,13 @@ def test_is_gemini_3_or_newer():
assert VertexGeminiConfig._is_gemini_3_or_newer("gemini-3-pro-preview") == True
assert VertexGeminiConfig._is_gemini_3_or_newer("gemini-3-flash") == True
assert VertexGeminiConfig._is_gemini_3_or_newer("gemini-3-pro") == True
- assert VertexGeminiConfig._is_gemini_3_or_newer("vertex_ai/gemini-3-pro-preview") == True
- assert VertexGeminiConfig._is_gemini_3_or_newer("gemini/gemini-3-pro-preview") == True
+ assert (
+ VertexGeminiConfig._is_gemini_3_or_newer("vertex_ai/gemini-3-pro-preview")
+ == True
+ )
+ assert (
+ VertexGeminiConfig._is_gemini_3_or_newer("gemini/gemini-3-pro-preview") == True
+ )
# Gemini 2.5 and older models
assert VertexGeminiConfig._is_gemini_3_or_newer("gemini-2.5-pro") == False
@@ -2209,8 +2276,12 @@ def test_media_resolution_from_detail_parameter():
)
# Test detail -> media_resolution enum mapping
- assert _convert_detail_to_media_resolution_enum("low") == {"level": "MEDIA_RESOLUTION_LOW"}
- assert _convert_detail_to_media_resolution_enum("high") == {"level": "MEDIA_RESOLUTION_HIGH"}
+ assert _convert_detail_to_media_resolution_enum("low") == {
+ "level": "MEDIA_RESOLUTION_LOW"
+ }
+ assert _convert_detail_to_media_resolution_enum("high") == {
+ "level": "MEDIA_RESOLUTION_HIGH"
+ }
assert _convert_detail_to_media_resolution_enum("auto") is None
assert _convert_detail_to_media_resolution_enum(None) is None
@@ -2223,19 +2294,16 @@ def test_media_resolution_from_detail_parameter():
"content": [
{
"type": "image_url",
- "image_url": {
- "url": base64_image,
- "detail": "high"
- }
+ "image_url": {"url": base64_image, "detail": "high"},
}
- ]
+ ],
}
]
contents = _gemini_convert_messages_with_history(
messages=messages, model="gemini-3-pro-preview"
)
-
+
# Verify media_resolution is set at the Part level (not inside inline_data)
assert len(contents) == 1
assert len(contents[0]["parts"]) >= 1
@@ -2266,19 +2334,16 @@ def test_media_resolution_low_detail():
"content": [
{
"type": "image_url",
- "image_url": {
- "url": base64_image,
- "detail": "low"
- }
+ "image_url": {"url": base64_image, "detail": "low"},
}
- ]
+ ],
}
]
contents = _gemini_convert_messages_with_history(
messages=messages, model="gemini-3-pro-preview"
)
-
+
# Find the part with inline_data
image_part = None
for part in contents[0]["parts"]:
@@ -2300,7 +2365,7 @@ def test_media_resolution_auto_detail():
# Using a minimal valid base64-encoded 1x1 PNG
base64_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
-
+
# Test with auto
messages_auto = [
{
@@ -2308,12 +2373,9 @@ def test_media_resolution_auto_detail():
"content": [
{
"type": "image_url",
- "image_url": {
- "url": base64_image,
- "detail": "auto"
- }
+ "image_url": {"url": base64_image, "detail": "auto"},
}
- ]
+ ],
}
]
@@ -2333,14 +2395,7 @@ def test_media_resolution_auto_detail():
messages_none = [
{
"role": "user",
- "content": [
- {
- "type": "image_url",
- "image_url": {
- "url": base64_image
- }
- }
- ]
+ "content": [{"type": "image_url", "image_url": {"url": base64_image}}],
}
]
@@ -2366,48 +2421,39 @@ def test_media_resolution_per_part():
# Using minimal valid base64-encoded 1x1 PNGs
base64_image1 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
base64_image2 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
-
+
messages = [
{
"role": "user",
"content": [
{
"type": "image_url",
- "image_url": {
- "url": base64_image1,
- "detail": "low"
- }
- },
- {
- "type": "text",
- "text": "Compare these images"
+ "image_url": {"url": base64_image1, "detail": "low"},
},
+ {"type": "text", "text": "Compare these images"},
{
"type": "image_url",
- "image_url": {
- "url": base64_image2,
- "detail": "high"
- }
- }
- ]
+ "image_url": {"url": base64_image2, "detail": "high"},
+ },
+ ],
}
]
contents = _gemini_convert_messages_with_history(
messages=messages, model="gemini-3-pro-preview"
)
-
+
# Should have one content with multiple parts
assert len(contents) == 1
assert len(contents[0]["parts"]) == 3 # image1, text, image2
-
+
# First image should have low resolution (first part is the image)
image1_part = contents[0]["parts"][0]
assert "inline_data" in image1_part
# media_resolution should be at the Part level, not inside inline_data
assert "media_resolution" in image1_part
assert image1_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_LOW"}
-
+
# Second image should have high resolution (third part is the second image)
image2_part = contents[0]["parts"][2]
assert "inline_data" in image2_part
@@ -2544,7 +2590,9 @@ def test_gemini_image_models_excluded_from_thinking():
)
# None of these should have thinkingConfig
- assert "thinkingConfig" not in result, f"Model {model} should not have thinkingConfig"
+ assert (
+ "thinkingConfig" not in result
+ ), f"Model {model} should not have thinkingConfig"
def test_partial_json_chunk_after_first_chunk():
@@ -2575,7 +2623,9 @@ def test_partial_json_chunk_after_first_chunk():
first_chunk = '{"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}'
result1 = iterator.handle_valid_json_chunk(first_chunk)
assert result1 is not None, "First complete chunk should parse OK"
- assert iterator.sent_first_chunk is True, "sent_first_chunk should be True after first chunk"
+ assert (
+ iterator.sent_first_chunk is True
+ ), "sent_first_chunk should be True after first chunk"
# Later chunk arrives PARTIAL (simulating network fragmentation)
partial_chunk = '{"candidates": [{"content":'
@@ -2583,7 +2633,9 @@ def test_partial_json_chunk_after_first_chunk():
# Should switch to accumulation mode instead of crashing
assert result2 is None, "Partial chunk should return None while accumulating"
- assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
+ assert (
+ iterator.chunk_type == "accumulated_json"
+ ), "Should switch to accumulated_json mode"
def test_partial_json_chunk_on_first_chunk():
@@ -2603,8 +2655,9 @@ def test_partial_json_chunk_on_first_chunk():
result = iterator.handle_valid_json_chunk(partial)
assert result is None, "Partial first chunk should return None"
- assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
-
+ assert (
+ iterator.chunk_type == "accumulated_json"
+ ), "Should switch to accumulated_json mode"
def test_google_ai_studio_presence_penalty_supported():
@@ -2617,6 +2670,8 @@ def test_google_ai_studio_presence_penalty_supported():
supported_params = config.get_supported_openai_params(model="gemini-2.0-flash")
assert "presence_penalty" in supported_params
+
+
# ==================== Tool Type Separation Tests ====================
# These tests verify that each Tool object contains exactly one type per Vertex AI API spec
# Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/Tool
@@ -2658,7 +2713,7 @@ def test_vertex_ai_multiple_tool_types_separate_objects():
{"enterpriseWebSearch": {}},
{"url_context": {}},
],
- optional_params=optional_params
+ optional_params=optional_params,
)
# Should have 2 separate Tool objects
@@ -2668,11 +2723,17 @@ def test_vertex_ai_multiple_tool_types_separate_objects():
tool_types_in_first = [k for k in tools[0].keys()]
tool_types_in_second = [k for k in tools[1].keys()]
- assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}"
- assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}"
+ assert (
+ len(tool_types_in_first) == 1
+ ), f"First Tool should have exactly 1 type, got {tool_types_in_first}"
+ assert (
+ len(tool_types_in_second) == 1
+ ), f"Second Tool should have exactly 1 type, got {tool_types_in_second}"
# Verify the correct tool types are present
- assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch"
+ assert (
+ "enterpriseWebSearch" in tools[0]
+ ), "First Tool should contain enterpriseWebSearch"
assert "url_context" in tools[1], "Second Tool should contain url_context"
@@ -2705,11 +2766,14 @@ def test_vertex_ai_function_declarations_with_other_tools_separate():
tools = v._map_function(
value=[
- {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}},
+ {
+ "type": "function",
+ "function": {"name": "get_weather", "description": "Get weather"},
+ },
{"googleSearch": {}},
{"code_execution": {}},
],
- optional_params=optional_params
+ optional_params=optional_params,
)
# Should have 2 Tool objects: function declarations + code_execution
@@ -2748,8 +2812,7 @@ def test_vertex_ai_single_tool_type_still_works():
optional_params = {}
tools = v._map_function(
- value=[{"code_execution": {}}],
- optional_params=optional_params
+ value=[{"code_execution": {}}], optional_params=optional_params
)
assert len(tools) == 1
@@ -2917,13 +2980,16 @@ def test_vertex_ai_openai_web_search_tool_transformation():
# Test web_search transformation
tools = v._map_function(
- value=[{"type": "web_search"}],
- optional_params=optional_params
+ value=[{"type": "web_search"}], optional_params=optional_params
)
assert len(tools) == 1, f"Expected 1 Tool object, got {len(tools)}"
- assert "googleSearch" in tools[0], f"Expected googleSearch in tool, got {tools[0].keys()}"
- assert tools[0]["googleSearch"] == {}, f"Expected empty googleSearch config, got {tools[0]['googleSearch']}"
+ assert (
+ "googleSearch" in tools[0]
+ ), f"Expected googleSearch in tool, got {tools[0].keys()}"
+ assert (
+ tools[0]["googleSearch"] == {}
+ ), f"Expected empty googleSearch config, got {tools[0]['googleSearch']}"
def test_vertex_ai_openai_web_search_preview_tool_transformation():
@@ -2941,13 +3007,16 @@ def test_vertex_ai_openai_web_search_preview_tool_transformation():
# Test web_search_preview transformation
tools = v._map_function(
- value=[{"type": "web_search_preview"}],
- optional_params=optional_params
+ value=[{"type": "web_search_preview"}], optional_params=optional_params
)
assert len(tools) == 1, f"Expected 1 Tool object, got {len(tools)}"
- assert "googleSearch" in tools[0], f"Expected googleSearch in tool, got {tools[0].keys()}"
- assert tools[0]["googleSearch"] == {}, f"Expected empty googleSearch config, got {tools[0]['googleSearch']}"
+ assert (
+ "googleSearch" in tools[0]
+ ), f"Expected googleSearch in tool, got {tools[0].keys()}"
+ assert (
+ tools[0]["googleSearch"] == {}
+ ), f"Expected empty googleSearch config, got {tools[0]['googleSearch']}"
def test_vertex_ai_openai_web_search_with_function_tools():
@@ -2973,9 +3042,12 @@ def test_vertex_ai_openai_web_search_with_function_tools():
tools = v._map_function(
value=[
{"type": "web_search"},
- {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}},
+ {
+ "type": "function",
+ "function": {"name": "get_weather", "description": "Get weather"},
+ },
],
- optional_params=optional_params
+ optional_params=optional_params,
)
# Should have 1 Tool object: function declarations only
@@ -3015,14 +3087,22 @@ def test_vertex_ai_multiple_function_declarations_grouped():
tools = v._map_function(
value=[
- {"type": "function", "function": {"name": "func1", "description": "First function"}},
- {"type": "function", "function": {"name": "func2", "description": "Second function"}},
+ {
+ "type": "function",
+ "function": {"name": "func1", "description": "First function"},
+ },
+ {
+ "type": "function",
+ "function": {"name": "func2", "description": "Second function"},
+ },
],
- optional_params=optional_params
+ optional_params=optional_params,
)
# Should have only 1 Tool object (function declarations grouped)
- assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}"
+ assert (
+ len(tools) == 1
+ ), f"Expected 1 Tool object for grouped functions, got {len(tools)}"
# Should contain function_declarations with 2 functions
assert "function_declarations" in tools[0]
@@ -3106,27 +3186,27 @@ def test_gemini_token_usage_standard_response():
def test_gemini_image_gen_usage_metadata_prompt_vs_completion_separation():
"""
Test that image generation models correctly separate prompt and completion token details.
-
+
This is a regression test for the bug where prompt_tokens_details.image_tokens
was incorrectly set to the completion's image token count instead of 0.
-
+
Scenario: Text-only prompt generates an image response
- Input: Text prompt (no images)
- Output: Generated image + text description
-
+
Expected behavior:
- prompt_tokens_details.image_tokens should be 0 (text-only input)
- completion_tokens_details.image_tokens should be 1290 (generated image)
-
+
Bug behavior (before fix):
- prompt_tokens_details.image_tokens was 1290 (incorrect!)
- completion_tokens_details.image_tokens was 1290 (correct)
-
+
The bug was caused by reusing the same variables (image_tokens, audio_tokens, text_tokens)
for both prompt and completion token details.
"""
v = VertexGeminiConfig()
-
+
# Simulate Gemini image generation model response metadata
# User sends text-only prompt, model generates image + text
usage_metadata_dict = {
@@ -3134,39 +3214,40 @@ def test_gemini_image_gen_usage_metadata_prompt_vs_completion_separation():
"candidatesTokenCount": 1290,
"totalTokenCount": 1391,
# Prompt is text-only (no image tokens in input)
- "promptTokensDetails": [
- {"modality": "TEXT", "tokenCount": 101}
- ],
+ "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 101}],
# Response contains generated image + text
- "candidatesTokensDetails": [
- {"modality": "IMAGE", "tokenCount": 1290}
- ],
+ "candidatesTokensDetails": [{"modality": "IMAGE", "tokenCount": 1290}],
}
-
+
completion_response = {"usageMetadata": usage_metadata_dict}
result = v._calculate_usage(completion_response=completion_response)
-
+
# Verify basic token counts
assert result.prompt_tokens == 101
assert result.completion_tokens == 1290
assert result.total_tokens == 1391
-
+
# CRITICAL: Prompt tokens details should show NO image tokens (text-only input)
- assert result.prompt_tokens_details.text_tokens == 101, \
- "Prompt text tokens should be 101"
- assert result.prompt_tokens_details.image_tokens is None, \
- "Prompt image tokens should be None (text-only input, no images in prompt)"
- assert result.prompt_tokens_details.audio_tokens is None, \
- "Prompt audio tokens should be None"
-
+ assert (
+ result.prompt_tokens_details.text_tokens == 101
+ ), "Prompt text tokens should be 101"
+ assert (
+ result.prompt_tokens_details.image_tokens is None
+ ), "Prompt image tokens should be None (text-only input, no images in prompt)"
+ assert (
+ result.prompt_tokens_details.audio_tokens is None
+ ), "Prompt audio tokens should be None"
+
# Completion tokens details should show the generated image tokens
- assert result.completion_tokens_details.image_tokens == 1290, \
- "Completion image tokens should be 1290 (generated image)"
-
+ assert (
+ result.completion_tokens_details.image_tokens == 1290
+ ), "Completion image tokens should be 1290 (generated image)"
+
# Verify text_tokens is auto-calculated for completion
# candidatesTokenCount (1290) - image_tokens (1290) = 0
- assert result.completion_tokens_details.text_tokens == 0, \
- "Completion text tokens should be 0 (image-only response)"
+ assert (
+ result.completion_tokens_details.text_tokens == 0
+ ), "Completion text tokens should be 0 (image-only response)"
def test_file_object_detail_parameter():
@@ -3185,10 +3266,10 @@ def test_file_object_detail_parameter():
"file": {
"file_id": "https://example.com/video.mp4",
"format": "video/mp4",
- "detail": "low"
- }
- }
- ]
+ "detail": "low",
+ },
+ },
+ ],
}
]
@@ -3208,7 +3289,9 @@ def test_file_object_detail_parameter():
break
assert file_part is not None, "File part should exist"
- assert "media_resolution" in file_part, "media_resolution should be set for file objects"
+ assert (
+ "media_resolution" in file_part
+ ), "media_resolution should be set for file objects"
assert file_part["media_resolution"] == {"level": "MEDIA_RESOLUTION_LOW"}
@@ -3228,10 +3311,10 @@ def test_video_metadata_fps():
"file": {
"file_id": "gs://bucket/video.mp4",
"format": "video/mp4",
- "video_metadata": {"fps": 5}
- }
- }
- ]
+ "video_metadata": {"fps": 5},
+ },
+ },
+ ],
}
]
@@ -3270,11 +3353,11 @@ def test_video_metadata_complete():
"video_metadata": {
"start_offset": "10s",
"end_offset": "60s",
- "fps": 5
- }
- }
- }
- ]
+ "fps": 5,
+ },
+ },
+ },
+ ],
}
]
@@ -3316,10 +3399,10 @@ def test_detail_and_video_metadata_combined():
"file_id": "https://example.com/video.mp4",
"format": "video/mp4",
"detail": "high",
- "video_metadata": {"fps": 10}
- }
- }
- ]
+ "video_metadata": {"fps": 10},
+ },
+ },
+ ],
}
]
@@ -3349,10 +3432,18 @@ def test_new_detail_levels():
)
# Test mapping function
- assert _convert_detail_to_media_resolution_enum("low") == {"level": "MEDIA_RESOLUTION_LOW"}
- assert _convert_detail_to_media_resolution_enum("medium") == {"level": "MEDIA_RESOLUTION_MEDIUM"}
- assert _convert_detail_to_media_resolution_enum("high") == {"level": "MEDIA_RESOLUTION_HIGH"}
- assert _convert_detail_to_media_resolution_enum("ultra_high") == {"level": "MEDIA_RESOLUTION_ULTRA_HIGH"}
+ assert _convert_detail_to_media_resolution_enum("low") == {
+ "level": "MEDIA_RESOLUTION_LOW"
+ }
+ assert _convert_detail_to_media_resolution_enum("medium") == {
+ "level": "MEDIA_RESOLUTION_MEDIUM"
+ }
+ assert _convert_detail_to_media_resolution_enum("high") == {
+ "level": "MEDIA_RESOLUTION_HIGH"
+ }
+ assert _convert_detail_to_media_resolution_enum("ultra_high") == {
+ "level": "MEDIA_RESOLUTION_ULTRA_HIGH"
+ }
# Test with actual message transformation
messages = [
@@ -3364,10 +3455,10 @@ def test_new_detail_levels():
"file": {
"file_id": "https://example.com/video.mp4",
"format": "video/mp4",
- "detail": "medium"
- }
+ "detail": "medium",
+ },
}
- ]
+ ],
}
]
@@ -3401,10 +3492,10 @@ def test_video_metadata_only_for_gemini_3():
"file_id": "https://example.com/video.mp4",
"format": "video/mp4",
"detail": "high",
- "video_metadata": {"fps": 5}
- }
+ "video_metadata": {"fps": 5},
+ },
}
- ]
+ ],
}
]
@@ -3420,8 +3511,12 @@ def test_video_metadata_only_for_gemini_3():
break
assert file_part_1_5 is not None
- assert "media_resolution" not in file_part_1_5, "Gemini 1.5 should not have media_resolution"
- assert "video_metadata" not in file_part_1_5, "Gemini 1.5 should not have video_metadata"
+ assert (
+ "media_resolution" not in file_part_1_5
+ ), "Gemini 1.5 should not have media_resolution"
+ assert (
+ "video_metadata" not in file_part_1_5
+ ), "Gemini 1.5 should not have video_metadata"
# Test with Gemini 3 (should have both)
contents_3 = _gemini_convert_messages_with_history(
@@ -3439,7 +3534,6 @@ def test_video_metadata_only_for_gemini_3():
assert "video_metadata" in file_part_3, "Gemini 3 should have video_metadata"
-
def test_chunk_parser_handles_prompt_feedback_block():
"""Test chunk_parser correctly handles promptFeedback.blockReason"""
from unittest.mock import Mock
@@ -3452,19 +3546,17 @@ def test_chunk_parser_handles_prompt_feedback_block():
blocked_chunk = {
"promptFeedback": {
"blockReason": "PROHIBITED_CONTENT",
- "blockReasonMessage": "The prompt is blocked due to prohibited contents"
+ "blockReasonMessage": "The prompt is blocked due to prohibited contents",
},
"responseId": "test_response_id",
- "modelVersion": "gemini-3-pro-preview"
+ "modelVersion": "gemini-3-pro-preview",
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
- streaming_response=iter([]),
- sync_stream=True,
- logging_obj=logging_obj
+ streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj
)
# Act
@@ -3473,7 +3565,9 @@ def test_chunk_parser_handles_prompt_feedback_block():
# Assert
assert result is not None, "Result should not be None"
assert len(result.choices) == 1, "Should have exactly one choice"
- assert result.choices[0].finish_reason == "content_filter", f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
+ assert (
+ result.choices[0].finish_reason == "content_filter"
+ ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
assert result.choices[0].delta.content is None, "content should be None"
@@ -3489,7 +3583,7 @@ def test_chunk_parser_handles_prompt_feedback_safety_block():
blocked_chunk = {
"promptFeedback": {
"blockReason": "SAFETY",
- "blockReasonMessage": "The prompt is blocked due to safety concerns"
+ "blockReasonMessage": "The prompt is blocked due to safety concerns",
},
"responseId": "test_safety_response_id",
}
@@ -3498,9 +3592,7 @@ def test_chunk_parser_handles_prompt_feedback_safety_block():
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
- streaming_response=iter([]),
- sync_stream=True,
- logging_obj=logging_obj
+ streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj
)
# Act
@@ -3524,24 +3616,22 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage():
blocked_chunk = {
"promptFeedback": {
"blockReason": "PROHIBITED_CONTENT",
- "blockReasonMessage": "The prompt is blocked due to prohibited contents"
+ "blockReasonMessage": "The prompt is blocked due to prohibited contents",
},
"responseId": "test_response_id_with_usage",
"modelVersion": "gemini-3-pro-preview",
"usageMetadata": {
"promptTokenCount": 8175,
"candidatesTokenCount": 0,
- "totalTokenCount": 8175
- }
+ "totalTokenCount": 8175,
+ },
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
- streaming_response=iter([]),
- sync_stream=True,
- logging_obj=logging_obj
+ streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj
)
# Act
@@ -3550,15 +3640,23 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage():
# Assert - 验证 content_filter 响应和 usage 都被正确处理
assert result is not None, "Result should not be None"
assert len(result.choices) == 1, "Should have exactly one choice"
- assert result.choices[0].finish_reason == "content_filter", f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
+ assert (
+ result.choices[0].finish_reason == "content_filter"
+ ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
assert result.choices[0].delta.content is None, "content should be None"
# 验证 usage 信息被正确提取
assert hasattr(result, "usage"), "result should have usage attribute"
assert result.usage is not None, "usage should not be None"
- assert result.usage.prompt_tokens == 8175, f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}"
- assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}"
- assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}"
+ assert (
+ result.usage.prompt_tokens == 8175
+ ), f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}"
+ assert (
+ result.usage.completion_tokens == 0
+ ), f"completion_tokens should be 0, got {result.usage.completion_tokens}"
+ assert (
+ result.usage.total_tokens == 8175
+ ), f"total_tokens should be 8175, got {result.usage.total_tokens}"
def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming():
@@ -3582,7 +3680,9 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming():
)
result = iterator.chunk_parser(chunk)
- assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND"
+ assert (
+ result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND"
+ )
def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming():
@@ -3621,7 +3721,10 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming():
encoding=None,
)
- assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "PROVISIONED_THROUGHPUT"
+ assert (
+ result._hidden_params["provider_specific_fields"]["traffic_type"]
+ == "PROVISIONED_THROUGHPUT"
+ )
def test_vertex_ai_service_tier_streaming():
@@ -3635,8 +3738,8 @@ def test_vertex_ai_service_tier_streaming():
}
iterator = ModelResponseIterator(
- streaming_response=[],
- sync_stream=True,
+ streaming_response=[],
+ sync_stream=True,
logging_obj=MagicMock(),
response_headers={"x-gemini-service-tier": "FLEX"},
)
@@ -3646,7 +3749,11 @@ def test_vertex_ai_service_tier_streaming():
# But definitely set when usageMetadata is present
chunk_with_usage = {
"candidates": [{"content": {"parts": [{"text": "hi"}]}}],
- "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}
+ "usageMetadata": {
+ "promptTokenCount": 1,
+ "candidatesTokenCount": 1,
+ "totalTokenCount": 2,
+ },
}
result_with_usage = iterator.chunk_parser(chunk_with_usage)
assert result_with_usage.service_tier == "flex"
@@ -3701,7 +3808,9 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api():
from litellm.types.utils import Choices, Message
model_response = ModelResponse()
- model_response._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND"}
+ model_response._hidden_params["provider_specific_fields"] = {
+ "traffic_type": "ON_DEMAND"
+ }
model_response.choices = [
Choices(
message=Message(content="Hello", role="assistant"),
@@ -3716,7 +3825,9 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api():
responses_api_request={},
)
- assert responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND"
+ assert (
+ responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND"
+ )
def test_vertex_ai_web_search_options_parameter():
@@ -3749,8 +3860,12 @@ def test_vertex_ai_web_search_options_parameter():
_tools = v._map_web_search_options(web_search_options)
# Verify the tool is a googleSearch tool
- assert "googleSearch" in _tools, f"Expected googleSearch in tool, got {_tools.keys()}"
- assert _tools["googleSearch"] == {}, f"Expected empty googleSearch config, got {_tools['googleSearch']}"
+ assert (
+ "googleSearch" in _tools
+ ), f"Expected googleSearch in tool, got {_tools.keys()}"
+ assert (
+ _tools["googleSearch"] == {}
+ ), f"Expected empty googleSearch config, got {_tools['googleSearch']}"
def test_vertex_ai_web_search_options_in_map_openai_params():
@@ -3773,14 +3888,14 @@ def test_vertex_ai_web_search_options_in_map_openai_params():
v = VertexGeminiConfig()
# Simulate optional_params passed to map_openai_params
- optional_params = {
- "web_search_options": {}
- }
+ optional_params = {"web_search_options": {}}
# Call the transformation that happens in map_openai_params
# Lines 1075-1079 in vertex_and_google_ai_studio_gemini.py (after fix)
web_search_value = optional_params.get("web_search_options")
- if isinstance(web_search_value, dict): # Fixed: removed 'value and' check to support empty dicts
+ if isinstance(
+ web_search_value, dict
+ ): # Fixed: removed 'value and' check to support empty dicts
_tools = v._map_web_search_options(web_search_value)
# Simulate _add_tools_to_optional_params
optional_params = v._add_tools_to_optional_params(optional_params, [_tools])
@@ -3792,8 +3907,12 @@ def test_vertex_ai_web_search_options_in_map_openai_params():
assert "tools" in optional_params, "tools should be added to optional_params"
assert len(optional_params["tools"]) == 1, "Should have exactly one tool"
assert "googleSearch" in optional_params["tools"][0], "Tool should be googleSearch"
- assert optional_params["tools"][0]["googleSearch"] == {}, "googleSearch should be empty config"
- assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation"
+ assert (
+ optional_params["tools"][0]["googleSearch"] == {}
+ ), "googleSearch should be empty config"
+ assert (
+ "web_search_options" not in optional_params
+ ), "web_search_options should be removed after transformation"
def test_vertex_ai_service_tier_in_map_openai_params():
@@ -3803,7 +3922,7 @@ def test_vertex_ai_service_tier_in_map_openai_params():
)
v = VertexGeminiConfig()
-
+
# Test pass-through
optional_params = {}
non_default_params = {"service_tier": "FLEX"}
@@ -3881,19 +4000,24 @@ def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt():
# Verify prompt token details include video tokens
assert result.prompt_tokens_details is not None
- assert result.prompt_tokens_details.video_tokens == 10240, \
- "Prompt video tokens should be 10240"
- assert result.prompt_tokens_details.text_tokens == 9, \
- "Prompt text tokens should be 9"
- assert result.prompt_tokens_details.audio_tokens == 200, \
- "Prompt audio tokens should be 200"
+ assert (
+ result.prompt_tokens_details.video_tokens == 10240
+ ), "Prompt video tokens should be 10240"
+ assert (
+ result.prompt_tokens_details.text_tokens == 9
+ ), "Prompt text tokens should be 9"
+ assert (
+ result.prompt_tokens_details.audio_tokens == 200
+ ), "Prompt audio tokens should be 200"
# Verify completion token details
assert result.completion_tokens_details is not None
- assert result.completion_tokens_details.text_tokens == 79, \
- "Completion text tokens should be 79"
- assert result.completion_tokens_details.video_tokens is None, \
- "Completion video tokens should be None (text-only response)"
+ assert (
+ result.completion_tokens_details.text_tokens == 79
+ ), "Completion text tokens should be 79"
+ assert (
+ result.completion_tokens_details.video_tokens is None
+ ), "Completion video tokens should be None (text-only response)"
def test_vertex_ai_usage_metadata_with_video_tokens_in_candidates():
@@ -3923,14 +4047,17 @@ def test_vertex_ai_usage_metadata_with_video_tokens_in_candidates():
assert result.completion_tokens == 10330
assert result.completion_tokens_details is not None
- assert result.completion_tokens_details.video_tokens == 10240, \
- "Completion video tokens should be 10240"
- assert result.completion_tokens_details.text_tokens == 90, \
- "Completion text tokens should be 90"
+ assert (
+ result.completion_tokens_details.video_tokens == 10240
+ ), "Completion video tokens should be 10240"
+ assert (
+ result.completion_tokens_details.text_tokens == 90
+ ), "Completion text tokens should be 90"
# Verify prompt side has no video tokens
- assert result.prompt_tokens_details.video_tokens is None, \
- "Prompt video tokens should be None (text-only input)"
+ assert (
+ result.prompt_tokens_details.video_tokens is None
+ ), "Prompt video tokens should be None (text-only input)"
def test_vertex_ai_usage_metadata_video_tokens_auto_calculated_text():
@@ -3956,8 +4083,9 @@ def test_vertex_ai_usage_metadata_video_tokens_auto_calculated_text():
assert result.completion_tokens_details.video_tokens == 10240
# text = 10330 - 10240 = 90
- assert result.completion_tokens_details.text_tokens == 90, \
- "text_tokens should be auto-calculated as candidatesTokenCount - video_tokens"
+ assert (
+ result.completion_tokens_details.text_tokens == 90
+ ), "text_tokens should be auto-calculated as candidatesTokenCount - video_tokens"
def test_vertex_ai_usage_metadata_video_tokens_with_caching():
@@ -3988,8 +4116,9 @@ def test_vertex_ai_usage_metadata_video_tokens_with_caching():
result = v._calculate_usage(completion_response=completion_response)
# video tokens should be reduced by cached amount: 10240 - 5120 = 5120
- assert result.prompt_tokens_details.video_tokens == 5120, \
- "Prompt video tokens should be 10240 - 5120 (cached) = 5120"
+ assert (
+ result.prompt_tokens_details.video_tokens == 5120
+ ), "Prompt video tokens should be 10240 - 5120 (cached) = 5120"
assert result.prompt_tokens_details.text_tokens == 9
assert result.prompt_tokens_details.audio_tokens == 200