From 29bf89cf9ce9ea7069dbeb4cca5f310ce018e1eb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 19 Jun 2025 13:54:02 -0700 Subject: [PATCH] =?UTF-8?q?fix(vertex=5Fai):=20Handle=20missing=20tokenCou?= =?UTF-8?q?nt=20in=20promptTokensDetails=20(#11=E2=80=A6=20(#11896)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(vertex_ai): Handle missing tokenCount in promptTokensDetails (#11581) This PR is a Solution to the Error converting to a valid response block='tokenCount'. File an issue if litellm error - https://github.com/BerriAI/litellm/issues It's happening because vertex_ai is not sometimes sending the token count for the audio modality. * test_vertex_ai_usage_metadata_missing_token_count --------- Co-authored-by: Nishith Jain <167524748+KingNish24@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 4 +-- ...test_vertex_and_google_ai_studio_gemini.py | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) 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 + +