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 91100cf7d7..e19bdc9edb 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 @@ -310,9 +310,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ return Tools(googleSearch={}) - def _transform_computer_use_config( - self, computer_use_config: dict - ) -> dict: + def _transform_computer_use_config(self, computer_use_config: dict) -> dict: """ Transform Computer Use configuration to Gemini API format. @@ -323,7 +321,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Transformed computer use configuration for Gemini API """ transformed_config = {} - + # Transform environment values if needed if "environment" in computer_use_config: env_value = computer_use_config["environment"] @@ -339,13 +337,17 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): f"Invalid environment value for computer_use: {env_value}. " f"Supported: 'browser', 'unspecified', 'ENVIRONMENT_BROWSER', 'ENVIRONMENT_UNSPECIFIED'" ) - + # Transform excluded_predefined_functions to camelCase if "excluded_predefined_functions" in computer_use_config: - transformed_config["excludedPredefinedFunctions"] = computer_use_config["excluded_predefined_functions"] + transformed_config["excludedPredefinedFunctions"] = computer_use_config[ + "excluded_predefined_functions" + ] elif "excludedPredefinedFunctions" in computer_use_config: - transformed_config["excludedPredefinedFunctions"] = computer_use_config["excludedPredefinedFunctions"] - + transformed_config["excludedPredefinedFunctions"] = computer_use_config[ + "excludedPredefinedFunctions" + ] + return transformed_config def _extract_google_maps_retrieval_config( @@ -446,9 +448,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"] @@ -553,7 +555,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "Invalid tool={}. Use `litellm.set_verbose` or `litellm --detailed_debug` to see raw request." ) -# Build list of Tool objects - each Tool should contain exactly one type + # Build list of Tool objects - each Tool should contain exactly one type # per Vertex AI API spec: "A Tool object should contain exactly one type of Tool" _tools_list: List[Tools] = [] @@ -570,11 +572,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() @@ -593,7 +599,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): computer_tool[VertexToolName.COMPUTER_USE.value] = computerUse _tools_list.append(computer_tool) - # Add retrieval config to toolConfig if googleMaps has location data if google_maps_retrieval_config is not None: if "toolConfig" not in optional_params: @@ -710,8 +715,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): GeminiThinkingConfig with thinkingLevel and includeThoughts """ # Check if this is gemini-3-flash which supports MINIMAL thinking level - is_gemini3flash= model and ( - "gemini-3-flash-preview" in model.lower() or "gemini-3-flash" in model.lower() + is_gemini3flash = model and ( + "gemini-3-flash-preview" in model.lower() + or "gemini-3-flash" in model.lower() ) if reasoning_effort == "minimal": if is_gemini3flash: @@ -799,7 +805,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): thinking_budget = thinking_param.get("budget_tokens") params: GeminiThinkingConfig = {} - + # For Gemini 3+ models, use thinkingLevel instead of thinkingBudget if model and VertexGeminiConfig._is_gemini_3_or_newer(model): if thinking_enabled: @@ -808,11 +814,21 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: params["includeThoughts"] = True if thinking_budget >= 10000: - is_gemini3flash = "gemini-3-flash-preview" in model.lower() or "gemini-3-flash" in model.lower() - params["thinkingLevel"] = "minimal" if is_gemini3flash else "low" + is_gemini3flash = ( + "gemini-3-flash-preview" in model.lower() + or "gemini-3-flash" in model.lower() + ) + params["thinkingLevel"] = ( + "minimal" if is_gemini3flash else "low" + ) else: - is_gemini3flash = "gemini-3-flash-preview" in model.lower() or "gemini-3-flash" in model.lower() - params["thinkingLevel"] = "minimal" if is_gemini3flash else "low" + is_gemini3flash = ( + "gemini-3-flash-preview" in model.lower() + or "gemini-3-flash" in model.lower() + ) + params["thinkingLevel"] = ( + "minimal" if is_gemini3flash else "low" + ) else: # Thinking disabled params["includeThoughts"] = False @@ -824,7 +840,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): params["includeThoughts"] = True if thinking_budget is not None and isinstance(thinking_budget, int): params["thinkingBudget"] = thinking_budget - + return params def map_response_modalities(self, value: list) -> list: @@ -980,16 +996,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( - value, model + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + value, model + ) ) else: - optional_params[ - "thinkingConfig" - ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - value, model + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( + value, model + ) ) elif param == "thinking": # Validate no conflict with thinking_level @@ -998,11 +1014,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) @@ -1036,8 +1052,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ): # For gemini-3-flash-preview, default to "minimal" to match Gemini 2.5 Flash behavior # For other Gemini 3 models, default to "low" - is_gemini3flash = "gemini-3-flash-preview" in model.lower() or "gemini-3-flash" in model.lower() - thinking_config["thinkingLevel"] = "minimal" if is_gemini3flash else "low" + is_gemini3flash = ( + "gemini-3-flash-preview" in model.lower() + or "gemini-3-flash" in model.lower() + ) + thinking_config["thinkingLevel"] = ( + "minimal" if is_gemini3flash else "low" + ) optional_params["thinkingConfig"] = thinking_config return optional_params @@ -1226,7 +1247,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): block: ChatCompletionThinkingBlock = { "type": "thinking", "thinking": thinking_text, - } + } signature = part.get("thoughtSignature") if signature is not None: block["signature"] = signature @@ -1360,10 +1381,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 @@ -1551,13 +1572,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif modality == "IMAGE": response_tokens_details.image_tokens = token_count - # Calculate text_tokens if not explicitly provided in candidatesTokensDetails - # candidatesTokenCount includes all modalities, so: text = total - (image + audio) + # Calculate text_tokens if not explicitly provided in candidatesTokensDetails + # candidatesTokenCount includes all modalities, so: text = total - (image + audio) + candidates_token_count = usage_metadata.get("candidatesTokenCount", 0) + if candidates_token_count > 0: + if response_tokens_details is None: + response_tokens_details = CompletionTokensDetailsWrapper() if response_tokens_details.text_tokens is None: - candidates_token_count = usage_metadata.get("candidatesTokenCount", 0) image_tokens = response_tokens_details.image_tokens or 0 audio_tokens_candidate = response_tokens_details.audio_tokens or 0 - calculated_text_tokens = candidates_token_count - image_tokens - audio_tokens_candidate + calculated_text_tokens = ( + candidates_token_count - image_tokens - audio_tokens_candidate + ) response_tokens_details.text_tokens = calculated_text_tokens ######################################################### @@ -2076,28 +2102,28 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## ADD METADATA TO RESPONSE ## setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) - model_response._hidden_params[ - "vertex_ai_grounding_metadata" - ] = grounding_metadata + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) setattr( model_response, "vertex_ai_url_context_metadata", url_context_metadata ) - model_response._hidden_params[ - "vertex_ai_url_context_metadata" - ] = url_context_metadata + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) setattr(model_response, "vertex_ai_safety_results", safety_ratings) - model_response._hidden_params[ - "vertex_ai_safety_results" - ] = safety_ratings # older approach - maintaining to prevent regressions + model_response._hidden_params["vertex_ai_safety_results"] = ( + safety_ratings # older approach - maintaining to prevent regressions + ) ## ADD CITATION METADATA ## setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) - model_response._hidden_params[ - "vertex_ai_citation_metadata" - ] = citation_metadata # older approach - maintaining to prevent regressions + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata # older approach - maintaining to prevent regressions + ) except Exception as e: raise VertexAIError( diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py new file mode 100644 index 0000000000..c8894ca7d5 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py @@ -0,0 +1,70 @@ + +import sys, os +import pytest +sys.path.insert(0, os.path.abspath('../../../../../')) + +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.types.llms.vertex_ai import UsageMetadata + +def test_gemini_3_flash_preview_token_usage_fallback(): + """Test fallback logic when candidatesTokensDetails is missing (e.g. Gemini 3 Flash Preview).""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 2145, + "candidatesTokenCount": 509, + "totalTokenCount": 2654, + # candidatesTokensDetails intentionally omitted + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 509 + assert result.prompt_tokens == 2145 + assert result.total_tokens == 2654 + + # Text tokens should be derived from candidatesTokenCount + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 509 + assert result.completion_tokens_details.image_tokens is None + assert result.completion_tokens_details.audio_tokens is None + +def test_gemini_no_reasoning_fallback(): + """Test fallback when reasoning_effort is absent and details are missing.""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 100, + "candidatesTokenCount": 264, + "totalTokenCount": 364, + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 264 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 264 + assert result.completion_tokens_details.reasoning_tokens is None or result.completion_tokens_details.reasoning_tokens == 0 + +def test_gemini_token_usage_standard_response(): + """Verify that standard responses with details are computed correctly and not overwritten.""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 40}, + {"modality": "IMAGE", "tokenCount": 10} + ] + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 50 + assert result.completion_tokens_details.text_tokens == 40 + assert result.completion_tokens_details.image_tokens == 10