From 24b768e27167f8d7e688232466e7602b8d21d2dd Mon Sep 17 00:00:00 2001 From: pradyyadav Date: Fri, 13 Mar 2026 01:02:07 +0530 Subject: [PATCH] address review: normalize timeout consistently, fix sync converse path, add tests Made-with: Cursor --- litellm/llms/bedrock/chat/converse_handler.py | 30 +- litellm/llms/bedrock/chat/invoke_handler.py | 12 +- .../vertex_and_google_ai_studio_gemini.py | 151 ++-- .../llms/bedrock/chat/test_invoke_handler.py | 22 +- ...test_vertex_and_google_ai_studio_gemini.py | 755 ++++++++---------- 5 files changed, 403 insertions(+), 567 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 4ab14d8b642..3e6f7e4be7b 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -42,7 +42,9 @@ 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, @@ -78,9 +80,7 @@ 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,10 +305,7 @@ 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( @@ -318,6 +315,7 @@ class BedrockConverseLLM(BaseAWSLLM): custom_llm_provider="bedrock", ) + ### SET REGION NAME ### aws_region_name = self._get_aws_region_name( optional_params=optional_params, @@ -375,7 +373,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" @@ -421,7 +419,7 @@ class BedrockConverseLLM(BaseAWSLLM): timeout=timeout, client=client, credentials=credentials, - api_key=api_key, + api_key=api_key ) # type: ignore ## TRANSFORMATION ## @@ -434,7 +432,7 @@ class BedrockConverseLLM(BaseAWSLLM): headers=extra_headers, ) data = json.dumps(_data) - + prepped = self.get_request_headers( credentials=credentials, aws_region_name=aws_region_name, @@ -442,7 +440,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 e7b0760ee95..518b21ac404 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -202,11 +202,7 @@ 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 @@ -417,9 +413,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 71a56e09a11..cc2fd24e906 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 @@ -498,9 +498,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"] @@ -510,19 +510,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "parameters" in _openai_function_object and _openai_function_object["parameters"] is not None and isinstance(_openai_function_object["parameters"], dict) - ): - 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"] - ) - else: - # Gemini 1.5: full OpenAPI-style transform - _openai_function_object["parameters"] = _build_vertex_schema( - _openai_function_object["parameters"] - ) + ): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema. + _openai_function_object["parameters"] = _build_vertex_schema( + _openai_function_object["parameters"] + ) openai_function_object = _openai_function_object @@ -641,15 +632,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() @@ -811,9 +802,12 @@ 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() + "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() ) - is_gemini31pro = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": if is_gemini3flash: return {"thinkingLevel": "minimal", "includeThoughts": True} @@ -1096,16 +1090,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 @@ -1114,11 +1108,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) @@ -1236,25 +1230,12 @@ 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]: @@ -1477,10 +1458,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 @@ -2290,37 +2271,35 @@ 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( @@ -2407,11 +2386,7 @@ 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() @@ -2990,11 +2965,7 @@ 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 @@ -3010,10 +2981,8 @@ 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, @@ -3046,9 +3015,7 @@ 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 4ca6bf53430..bbd63c6a9eb 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -249,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(): @@ -311,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(): @@ -338,7 +338,9 @@ 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", @@ -351,6 +353,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 cdea6981be1..94c0ddad7b4 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,9 +237,7 @@ 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(): @@ -319,22 +317,14 @@ 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(): @@ -591,7 +581,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", @@ -610,18 +600,12 @@ 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(): @@ -669,15 +653,12 @@ 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(): @@ -730,10 +711,7 @@ 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) @@ -763,9 +741,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}) @@ -807,7 +785,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}) @@ -831,13 +809,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) """ @@ -848,29 +826,31 @@ 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 ( @@ -900,17 +880,13 @@ 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"] == {} @@ -1076,13 +1052,7 @@ 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"]} ], } ], @@ -1426,7 +1396,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. """ @@ -1443,7 +1413,12 @@ 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( @@ -1458,27 +1433,19 @@ 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" @@ -1493,17 +1460,15 @@ 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 @@ -1513,49 +1478,45 @@ 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" @@ -1566,7 +1527,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": { @@ -1577,7 +1538,7 @@ def test_vertex_ai_map_google_maps_tool_with_location(): } }] optional_params={} - + Expected Output: tools=[{ "googleMaps": {"enableWidget": "ENABLE_WIDGET"} @@ -1596,43 +1557,40 @@ 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 @@ -1647,19 +1605,14 @@ 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" @@ -1667,7 +1620,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 = {} @@ -1675,16 +1628,12 @@ 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" @@ -1696,7 +1645,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 @@ -1713,25 +1662,22 @@ 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 = {} @@ -1739,46 +1685,39 @@ 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 @@ -1801,7 +1740,7 @@ def test_vertex_ai_annotation_streaming_events(): "groundingMetadata": { "webSearchQueries": ["weather San Francisco today"], "searchEntryPoint": { - "renderedContent": "
Search results
" + "renderedContent": '
Search results
' }, "groundingChunks": [ { @@ -1842,7 +1781,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") @@ -1895,7 +1834,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. """ @@ -1906,7 +1845,9 @@ 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": { @@ -1921,7 +1862,7 @@ def test_vertex_ai_annotation_conversion(): "title": "Current time in San Francisco, CA", "domain": "google.com", } - }, + } ], "groundingSupports": [ { @@ -1950,14 +1891,12 @@ 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 ) @@ -1993,7 +1932,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 ( @@ -2031,7 +1970,6 @@ 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 ( @@ -2042,13 +1980,8 @@ 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 @@ -2240,12 +2173,8 @@ 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 @@ -2258,16 +2187,19 @@ 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 @@ -2298,16 +2230,19 @@ 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"]: @@ -2329,7 +2264,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 = [ { @@ -2337,9 +2272,12 @@ def test_media_resolution_auto_detail(): "content": [ { "type": "image_url", - "image_url": {"url": base64_image, "detail": "auto"}, + "image_url": { + "url": base64_image, + "detail": "auto" + } } - ], + ] } ] @@ -2359,7 +2297,14 @@ 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 + } + } + ] } ] @@ -2385,39 +2330,48 @@ 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"}, + "image_url": { + "url": base64_image1, + "detail": "low" + } + }, + { + "type": "text", + "text": "Compare these images" }, - {"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 @@ -2554,9 +2508,7 @@ 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(): @@ -2587,9 +2539,7 @@ 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":' @@ -2597,9 +2547,7 @@ 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(): @@ -2619,9 +2567,8 @@ 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(): @@ -2634,8 +2581,6 @@ 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 @@ -2677,7 +2622,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 @@ -2687,17 +2632,11 @@ 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" @@ -2727,14 +2666,11 @@ 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 @@ -2759,9 +2695,7 @@ 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" @@ -2783,7 +2717,8 @@ 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 @@ -2812,16 +2747,13 @@ 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(): @@ -2839,16 +2771,13 @@ 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(): @@ -2873,12 +2802,9 @@ 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 @@ -2933,22 +2859,14 @@ 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] @@ -3032,27 +2950,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 = { @@ -3060,40 +2978,39 @@ 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(): @@ -3112,10 +3029,10 @@ def test_file_object_detail_parameter(): "file": { "file_id": "https://example.com/video.mp4", "format": "video/mp4", - "detail": "low", - }, - }, - ], + "detail": "low" + } + } + ] } ] @@ -3135,9 +3052,7 @@ 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"} @@ -3157,10 +3072,10 @@ def test_video_metadata_fps(): "file": { "file_id": "gs://bucket/video.mp4", "format": "video/mp4", - "video_metadata": {"fps": 5}, - }, - }, - ], + "video_metadata": {"fps": 5} + } + } + ] } ] @@ -3199,11 +3114,11 @@ def test_video_metadata_complete(): "video_metadata": { "start_offset": "10s", "end_offset": "60s", - "fps": 5, - }, - }, - }, - ], + "fps": 5 + } + } + } + ] } ] @@ -3245,10 +3160,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} + } + } + ] } ] @@ -3278,18 +3193,10 @@ 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 = [ @@ -3301,10 +3208,10 @@ def test_new_detail_levels(): "file": { "file_id": "https://example.com/video.mp4", "format": "video/mp4", - "detail": "medium", - }, + "detail": "medium" + } } - ], + ] } ] @@ -3338,10 +3245,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} + } } - ], + ] } ] @@ -3357,12 +3264,8 @@ 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( @@ -3380,6 +3283,7 @@ 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 @@ -3392,17 +3296,19 @@ 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 @@ -3411,9 +3317,7 @@ 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" @@ -3429,7 +3333,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", } @@ -3438,7 +3342,9 @@ 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 @@ -3462,22 +3368,24 @@ 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 @@ -3486,23 +3394,15 @@ 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(): @@ -3526,9 +3426,7 @@ 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(): @@ -3567,10 +3465,7 @@ 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(): @@ -3583,9 +3478,7 @@ 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"), @@ -3600,9 +3493,7 @@ 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(): @@ -3635,12 +3526,8 @@ 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(): @@ -3663,14 +3550,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]) @@ -3682,12 +3569,8 @@ 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(): @@ -3727,24 +3610,19 @@ 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(): @@ -3774,17 +3652,14 @@ 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(): @@ -3810,9 +3685,8 @@ 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(): @@ -3843,9 +3717,8 @@ 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 @@ -3960,9 +3833,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(): @@ -4012,9 +3885,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(): @@ -4059,6 +3932,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" + )