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 e20236e772..e058aa675d 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 @@ -1025,9 +1025,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): response_tokens_details = CompletionTokensDetailsWrapper() for detail in usage_metadata["responseTokensDetails"]: if detail["modality"] == "TEXT": - response_tokens_details.text_tokens = detail["tokenCount"] + response_tokens_details.text_tokens = detail.get("tokenCount", 0) elif detail["modality"] == "AUDIO": - response_tokens_details.audio_tokens = detail["tokenCount"] + response_tokens_details.audio_tokens = detail.get("tokenCount", 0) ######################################################### if "promptTokensDetails" in usage_metadata: 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 e47cac75dd..c7c7b2cb66 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 @@ -713,3 +713,31 @@ def test_vertex_ai_transform_parts(): assert function["name"] == "simple_function" assert function["arguments"] == "{}" assert tools is None + + +def test_vertex_ai_usage_metadata_missing_token_count(): + """Test that missing tokenCount in responseTokensDetails defaults to 0""" + from litellm.types.utils import PromptTokensDetailsWrapper + + v = VertexGeminiConfig() + usage_metadata = { + "promptTokenCount": 57, + "responseTokenCount": 74, + "totalTokenCount": 131, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 57}], + "responseTokensDetails": [ + {"modality": "TEXT"}, # Missing tokenCount + {"modality": "AUDIO"}, # Missing tokenCount + ], + } + usage_metadata = UsageMetadata(**usage_metadata) + result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata}) + + # Should not crash and should default missing tokenCount to 0 + assert result.prompt_tokens == 57 + assert result.completion_tokens == 74 + assert result.total_tokens == 131 + assert result.completion_tokens_details.text_tokens == 0 # Default value for missing tokenCount + assert result.completion_tokens_details.audio_tokens == 0 # Default value for missing tokenCount + +