From 2847437c04d1075e7a68118c0f3b379ce6137884 Mon Sep 17 00:00:00 2001 From: pradyyadav Date: Fri, 13 Mar 2026 00:38:04 +0530 Subject: [PATCH] address review: normalize timeout consistently, fix sync converse path, add tests --- litellm/llms/bedrock/chat/converse_handler.py | 30 +- litellm/llms/bedrock/chat/invoke_handler.py | 12 +- .../vertex_and_google_ai_studio_gemini.py | 148 ++-- .../llms/bedrock/chat/test_invoke_handler.py | 29 +- ...test_vertex_and_google_ai_studio_gemini.py | 756 ++++++++++-------- 5 files changed, 566 insertions(+), 409 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 3e6f7e4be7b..4ab14d8b642 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -42,9 +42,7 @@ def make_sync_call( _params: dict = {} if timeout is not None: _params["timeout"] = timeout - client = _get_httpx_client( - params=_params if _params else None - ) + client = _get_httpx_client(params=_params if _params else None) response = client.post( api_base, @@ -80,7 +78,9 @@ def make_sync_call( ) else: decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) + completion_stream = decoder.iter_bytes( + response.iter_bytes(chunk_size=stream_chunk_size) + ) # LOGGING logging_obj.post_call( @@ -134,7 +134,7 @@ class BedrockConverseLLM(BaseAWSLLM): endpoint_url=api_base, data=data, headers=headers, - api_key=api_key + api_key=api_key, ) ## LOGGING @@ -195,7 +195,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers=headers, ) data = json.dumps(request_data) - + prepped = self.get_request_headers( credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", @@ -203,7 +203,7 @@ class BedrockConverseLLM(BaseAWSLLM): endpoint_url=api_base, data=data, headers=headers, - api_key=api_key + api_key=api_key, ) ## LOGGING @@ -289,7 +289,7 @@ class BedrockConverseLLM(BaseAWSLLM): _stripped = _model_for_id for rp in ["bedrock/converse/", "bedrock/", "converse/"]: if _stripped.startswith(rp): - _stripped = _stripped[len(rp):] + _stripped = _stripped[len(rp) :] break # Strip embedded region prefix (e.g. "bedrock/us-east-1/model" -> "model") # and capture it so it can be used as aws_region_name below. @@ -305,7 +305,10 @@ class BedrockConverseLLM(BaseAWSLLM): break modelId = self.encode_model_id(model_id=_model_for_id) # Inject region extracted from model path so _get_aws_region_name picks it up - if _region_from_model is not None and "aws_region_name" not in optional_params: + if ( + _region_from_model is not None + and "aws_region_name" not in optional_params + ): optional_params["aws_region_name"] = _region_from_model fake_stream = litellm.AmazonConverseConfig().should_fake_stream( @@ -315,7 +318,6 @@ class BedrockConverseLLM(BaseAWSLLM): custom_llm_provider="bedrock", ) - ### SET REGION NAME ### aws_region_name = self._get_aws_region_name( optional_params=optional_params, @@ -373,7 +375,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - + # Filter beta headers in HTTP headers before making the request headers = update_headers_with_filtered_beta( headers=headers, provider="bedrock_converse" @@ -419,7 +421,7 @@ class BedrockConverseLLM(BaseAWSLLM): timeout=timeout, client=client, credentials=credentials, - api_key=api_key + api_key=api_key, ) # type: ignore ## TRANSFORMATION ## @@ -432,7 +434,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers=extra_headers, ) data = json.dumps(_data) - + prepped = self.get_request_headers( credentials=credentials, aws_region_name=aws_region_name, @@ -440,7 +442,7 @@ class BedrockConverseLLM(BaseAWSLLM): endpoint_url=proxy_endpoint_url, data=data, headers=headers, - api_key=api_key + api_key=api_key, ) ## LOGGING diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 518b21ac404..e7b0760ee95 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -202,7 +202,11 @@ async def make_call( if client is None: _params: dict = {} - if logging_obj and logging_obj.litellm_params and logging_obj.litellm_params.get("ssl_verify"): + if ( + logging_obj + and logging_obj.litellm_params + and logging_obj.litellm_params.get("ssl_verify") + ): _params["ssl_verify"] = logging_obj.litellm_params.get("ssl_verify") if timeout is not None: _params["timeout"] = timeout @@ -413,9 +417,9 @@ class BedrockLLM(BaseAWSLLM): # Claude 3+ indicators (all use Messages API) messages_api_indicators = [ - "claude-3", # Claude 3.x models - "claude-opus-4", # Claude Opus 4 - "claude-sonnet-4", # Claude Sonnet 4 + "claude-3", # Claude 3.x models + "claude-opus-4", # Claude Opus 4 + "claude-sonnet-4", # Claude Sonnet 4 "claude-haiku-4", # Claude Haiku 4 ] 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 1258e2b72bd..95705c0b75a 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 @@ -499,9 +499,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"] @@ -514,17 +514,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ): if supports_response_json_schema(model): # Gemini 2.0+: minimal transform (resolve $ref only) - _openai_function_object["parameters"] = ( - _build_vertex_schema_for_gemini_2( - _openai_function_object["parameters"] - ) + _openai_function_object[ + "parameters" + ] = _build_vertex_schema_for_gemini_2( + _openai_function_object["parameters"] ) else: # Gemini 1.5: full OpenAPI-style transform - _openai_function_object["parameters"] = ( - _build_vertex_schema( - _openai_function_object["parameters"] - ) + _openai_function_object["parameters"] = _build_vertex_schema( + _openai_function_object["parameters"] ) openai_function_object = _openai_function_object @@ -644,15 +642,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() @@ -814,12 +812,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Check if this is gemini-3-flash which supports MINIMAL thinking level # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, etc. is_gemini3flash = model and ( - "gemini-3-flash" in model.lower() - or "gemini-3.1-flash" in model.lower() - ) - is_gemini31pro = model and ( - "gemini-3.1-pro-preview" in model.lower() + "gemini-3-flash" in model.lower() or "gemini-3.1-flash" in model.lower() ) + is_gemini31pro = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": if is_gemini3flash: return {"thinkingLevel": "minimal", "includeThoughts": True} @@ -1102,16 +1097,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 @@ -1120,11 +1115,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) @@ -1242,12 +1237,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } - _GEMINI_FINISH_REASON_KEYS = frozenset({ - "STOP", "MAX_TOKENS", "SAFETY", "RECITATION", "FINISH_REASON_UNSPECIFIED", - "MALFORMED_FUNCTION_CALL", "LANGUAGE", "OTHER", "BLOCKLIST", - "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", - "TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE", - }) + _GEMINI_FINISH_REASON_KEYS = frozenset( + { + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "FINISH_REASON_UNSPECIFIED", + "MALFORMED_FUNCTION_CALL", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "TOO_MANY_TOOL_CALLS", + "MALFORMED_RESPONSE", + } + ) @staticmethod def get_finish_reason_mapping() -> Dict[str, OpenAIChatCompletionFinishReason]: @@ -1470,10 +1478,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 @@ -2283,35 +2291,37 @@ 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( "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 except Exception as e: raise VertexAIError( @@ -2398,7 +2408,11 @@ async def make_call( try: response = await client.post( - api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, + api_base, + headers=headers, + data=data, + stream=True, + logging_obj=logging_obj, timeout=timeout, ) response.raise_for_status() @@ -2977,7 +2991,11 @@ class ModelResponseIterator: # to correctly set finish_reason="tool_calls" per the OpenAI spec. if not self.has_seen_tool_calls: for choice in model_response.choices: - if hasattr(choice, "delta") and choice.delta and choice.delta.tool_calls: + if ( + hasattr(choice, "delta") + and choice.delta + and choice.delta.tool_calls + ): self.has_seen_tool_calls = True break @@ -2993,8 +3011,10 @@ class ModelResponseIterator: if self.has_seen_tool_calls: mapped_finish_reason = "tool_calls" else: - mapped_finish_reason = VertexGeminiConfig._check_finish_reason( - None, finish_reason_str + mapped_finish_reason = ( + VertexGeminiConfig._check_finish_reason( + None, finish_reason_str + ) ) choice = StreamingChoices( finish_reason=mapped_finish_reason, @@ -3027,7 +3047,9 @@ class ModelResponseIterator: "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 setattr(model_response, "usage", usage) # type: ignore diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index f178d8b69ae..4ca6bf53430 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -211,7 +211,7 @@ def test_bedrock_invoke_async_streaming_passes_timeout_to_make_call(): Fixes https://github.com/BerriAI/litellm/issues/23375 """ import asyncio - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import MagicMock, patch import httpx @@ -229,6 +229,7 @@ def test_bedrock_invoke_async_streaming_passes_timeout_to_make_call(): "litellm.llms.bedrock.chat.invoke_handler.CustomStreamWrapper", FakeCustomStreamWrapper, ): + async def run(): await handler.async_streaming( model="anthropic.claude-3-sonnet", @@ -248,9 +249,9 @@ def test_bedrock_invoke_async_streaming_passes_timeout_to_make_call(): make_call_partial = captured_partial.get("make_call") assert make_call_partial is not None - assert make_call_partial.keywords.get("timeout") == timeout, ( - "timeout must be forwarded via partial() to make_call()" - ) + assert ( + make_call_partial.keywords.get("timeout") == timeout + ), "timeout must be forwarded via partial() to make_call()" def test_bedrock_converse_async_streaming_passes_timeout_to_make_call(): @@ -284,7 +285,9 @@ def test_bedrock_converse_async_streaming_passes_timeout_to_make_call(): ) as mock_make_call, patch( "litellm.AmazonConverseConfig", ) as mock_converse_config, patch.object( - handler, "get_request_headers", return_value=fake_prepped, + handler, + "get_request_headers", + return_value=fake_prepped, ): mock_converse_config.return_value._async_transform_request = AsyncMock( return_value={"messages": []} @@ -308,9 +311,9 @@ def test_bedrock_converse_async_streaming_passes_timeout_to_make_call(): asyncio.run(run()) _, kwargs = mock_make_call.call_args - assert kwargs.get("timeout") == timeout, ( - "timeout must be forwarded to make_call() in BedrockConverseLLM.async_streaming()" - ) + assert ( + kwargs.get("timeout") == timeout + ), "timeout must be forwarded to make_call() in BedrockConverseLLM.async_streaming()" def test_bedrock_converse_sync_make_sync_call_passes_timeout_to_client_post(): @@ -335,9 +338,7 @@ def test_bedrock_converse_sync_make_sync_call_passes_timeout_to_client_post(): timeout = httpx.Timeout(4.0) - with patch( - "litellm.llms.bedrock.chat.converse_handler.AWSEventStreamDecoder" - ): + with patch("litellm.llms.bedrock.chat.converse_handler.AWSEventStreamDecoder"): make_sync_call( client=mock_client, api_base="https://example.com", @@ -350,6 +351,6 @@ def test_bedrock_converse_sync_make_sync_call_passes_timeout_to_client_post(): ) _, kwargs = mock_client.post.call_args - assert kwargs.get("timeout") == timeout, ( - "timeout must be forwarded to client.post() in converse make_sync_call()" - ) + assert ( + kwargs.get("timeout") == timeout + ), "timeout must be forwarded to client.post() in converse make_sync_call()" 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 ecc9d645ab2..cdea6981be1 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 ( @@ -880,13 +900,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"] == {} @@ -1052,7 +1076,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", + ] + } ], } ], @@ -1396,7 +1426,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. """ @@ -1413,12 +1443,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( @@ -1433,19 +1458,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" @@ -1460,15 +1493,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 @@ -1478,45 +1513,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" @@ -1527,7 +1566,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": { @@ -1538,7 +1577,7 @@ def test_vertex_ai_map_google_maps_tool_with_location(): } }] optional_params={} - + Expected Output: tools=[{ "googleMaps": {"enableWidget": "ENABLE_WIDGET"} @@ -1557,40 +1596,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 @@ -1605,14 +1647,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" @@ -1620,7 +1667,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 = {} @@ -1628,12 +1675,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" @@ -1645,7 +1696,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 @@ -1662,22 +1713,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 = {} @@ -1685,39 +1739,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 @@ -1740,7 +1801,7 @@ def test_vertex_ai_annotation_streaming_events(): "groundingMetadata": { "webSearchQueries": ["weather San Francisco today"], "searchEntryPoint": { - "renderedContent": '
Search results
' + "renderedContent": "
Search results
" }, "groundingChunks": [ { @@ -1781,7 +1842,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") @@ -1834,7 +1895,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. """ @@ -1845,9 +1906,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": { @@ -1862,7 +1921,7 @@ def test_vertex_ai_annotation_conversion(): "title": "Current time in San Francisco, CA", "domain": "google.com", } - } + }, ], "groundingSupports": [ { @@ -1891,12 +1950,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 ) @@ -1932,7 +1993,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 ( @@ -1970,6 +2031,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 ( @@ -1980,8 +2042,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 @@ -2173,8 +2240,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 @@ -2187,19 +2258,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 @@ -2230,19 +2298,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"]: @@ -2264,7 +2329,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 = [ { @@ -2272,12 +2337,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"}, } - ] + ], } ] @@ -2297,14 +2359,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}}], } ] @@ -2330,48 +2385,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 @@ -2508,7 +2554,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(): @@ -2539,7 +2587,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":' @@ -2547,7 +2597,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(): @@ -2567,8 +2619,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(): @@ -2581,6 +2634,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 @@ -2622,7 +2677,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 @@ -2632,11 +2687,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" @@ -2666,11 +2727,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 3 separate Tool objects @@ -2695,7 +2759,9 @@ def test_vertex_ai_function_declarations_with_other_tools_separate(): assert code_tool is not None, "code_execution Tool should be present" # Verify each Tool has exactly one type - assert len(func_tool.keys()) == 1, "function_declarations Tool should have only one key" + assert ( + len(func_tool.keys()) == 1 + ), "function_declarations Tool should have only one key" assert len(search_tool.keys()) == 1, "googleSearch Tool should have only one key" assert len(code_tool.keys()) == 1, "code_execution Tool should have only one key" @@ -2717,8 +2783,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 @@ -2747,13 +2812,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(): @@ -2771,13 +2839,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(): @@ -2802,9 +2873,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 2 separate Tool objects @@ -2859,14 +2933,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] @@ -2950,27 +3032,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 = { @@ -2978,39 +3060,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(): @@ -3029,10 +3112,10 @@ def test_file_object_detail_parameter(): "file": { "file_id": "https://example.com/video.mp4", "format": "video/mp4", - "detail": "low" - } - } - ] + "detail": "low", + }, + }, + ], } ] @@ -3052,7 +3135,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"} @@ -3072,10 +3157,10 @@ def test_video_metadata_fps(): "file": { "file_id": "gs://bucket/video.mp4", "format": "video/mp4", - "video_metadata": {"fps": 5} - } - } - ] + "video_metadata": {"fps": 5}, + }, + }, + ], } ] @@ -3114,11 +3199,11 @@ def test_video_metadata_complete(): "video_metadata": { "start_offset": "10s", "end_offset": "60s", - "fps": 5 - } - } - } - ] + "fps": 5, + }, + }, + }, + ], } ] @@ -3160,10 +3245,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}, + }, + }, + ], } ] @@ -3193,10 +3278,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 = [ @@ -3208,10 +3301,10 @@ def test_new_detail_levels(): "file": { "file_id": "https://example.com/video.mp4", "format": "video/mp4", - "detail": "medium" - } + "detail": "medium", + }, } - ] + ], } ] @@ -3245,10 +3338,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}, + }, } - ] + ], } ] @@ -3264,8 +3357,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( @@ -3283,7 +3380,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 @@ -3296,19 +3392,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 @@ -3317,7 +3411,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" @@ -3333,7 +3429,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", } @@ -3342,9 +3438,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 @@ -3368,24 +3462,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 @@ -3394,15 +3486,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(): @@ -3426,7 +3526,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(): @@ -3465,7 +3567,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_traffic_type_surfaced_in_responses_api(): @@ -3478,7 +3583,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"), @@ -3493,7 +3600,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(): @@ -3526,8 +3635,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(): @@ -3550,14 +3663,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]) @@ -3569,8 +3682,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_usage_metadata_with_video_tokens_in_prompt(): @@ -3610,19 +3727,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(): @@ -3652,14 +3774,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(): @@ -3685,8 +3810,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(): @@ -3717,8 +3843,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 @@ -3833,9 +3960,9 @@ def test_vertex_make_call_passes_timeout_to_client_post(): asyncio.run(run()) _, kwargs = mock_client.post.call_args - assert kwargs.get("timeout") == timeout, ( - "timeout must be forwarded to client.post() for streaming to respect it" - ) + assert ( + kwargs.get("timeout") == timeout + ), "timeout must be forwarded to client.post() for streaming to respect it" def test_vertex_make_call_creates_client_with_timeout_when_no_client_provided(): @@ -3868,6 +3995,7 @@ def test_vertex_make_call_creates_client_with_timeout_when_no_client_provided(): "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.get_async_httpx_client", return_value=mock_created_client, ) as mock_get_client: + async def run(): await make_call( client=None, @@ -3884,9 +4012,9 @@ def test_vertex_make_call_creates_client_with_timeout_when_no_client_provided(): asyncio.run(run()) _, kwargs = mock_get_client.call_args - assert kwargs.get("params", {}).get("timeout") == timeout, ( - "timeout must be passed to get_async_httpx_client() params" - ) + assert ( + kwargs.get("params", {}).get("timeout") == timeout + ), "timeout must be passed to get_async_httpx_client() params" def test_vertex_make_call_passes_timeout_with_gemini_client(): @@ -3931,6 +4059,6 @@ def test_vertex_make_call_passes_timeout_with_gemini_client(): asyncio.run(run()) _, kwargs = mock_gemini_client.post.call_args - assert kwargs.get("timeout") == timeout, ( - "timeout must be forwarded to client.post() when gemini_client is provided" - ) + assert ( + kwargs.get("timeout") == timeout + ), "timeout must be forwarded to client.post() when gemini_client is provided"