ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fixes #18896 : Handle missing completion_tokens_details when reasoning_effort is not used

Browse Source
This commit is contained in:
yogeshwaran10 2026-01-11 00:45:42 +05:30
parent 3a2de85e7b
commit d98c71f07e
2 changed files with 155 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View File

@ -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(

70
tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py Normal file
View File

@ -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