fix(vertex_ai): Handle missing tokenCount in promptTokensDetails (#11… (#11896)

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

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:

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