fix(gemini): include DOCUMENT modality tokens in cost calculation

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -1721,6 +1721,10 @@ def _get_token_count(detail: Mapping[str, Any]) -> int:
                     response_tokens_details.audio_tokens = (
                         response_tokens_details.audio_tokens or 0
                     ) + token_count
+                elif modality == "DOCUMENT":
+                    response_tokens_details.text_tokens = (
+                        response_tokens_details.text_tokens or 0
+                    ) + token_count
 
         #########################################################
 
@@ -1747,6 +1751,10 @@ def _get_token_count(detail: Mapping[str, Any]) -> int:
                     response_tokens_details.video_tokens = (
                         response_tokens_details.video_tokens or 0
                     ) + token_count
+                elif modality == "DOCUMENT":
+                    response_tokens_details.text_tokens = (
+                        response_tokens_details.text_tokens or 0
+                    ) + token_count
 
         # Calculate text_tokens if not explicitly provided in candidatesTokensDetails
         # candidatesTokenCount includes all modalities, so: text = total - (image + audio + video)
@@ -1780,6 +1788,8 @@ def _get_token_count(detail: Mapping[str, Any]) -> int:
                     prompt_image_tokens = (prompt_image_tokens or 0) + token_count
                 elif modality == "VIDEO":
                     prompt_video_tokens = (prompt_video_tokens or 0) + token_count
+                elif modality == "DOCUMENT":
+                    prompt_text_tokens = (prompt_text_tokens or 0) + token_count
 
         ## Parse cacheTokensDetails (breakdown of cached tokens by modality)
         ## When explicit caching is used, Gemini provides this field to show which modalities were cached
@@ -1800,6 +1810,8 @@ def _get_token_count(detail: Mapping[str, Any]) -> int:
                     cached_image_tokens = (cached_image_tokens or 0) + token_count
                 elif modality == "VIDEO":
                     cached_video_tokens = (cached_video_tokens or 0) + token_count
+                elif modality == "DOCUMENT":
+                    cached_text_tokens = (cached_text_tokens or 0) + token_count
 
         ## Calculate non-cached tokens by subtracting cached from total (per modality)
         ## This is necessary because promptTokensDetails includes both cached and non-cached tokens

tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py

@@ -3759,6 +3759,88 @@ def test_vertex_ai_usage_metadata_video_tokens_with_caching():
     assert result.prompt_tokens_details.audio_tokens == 200
 
 
+def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt():
+    """Test promptTokensDetails with DOCUMENT modality for PDF inputs.
+
+    Fixes: https://github.com/BerriAI/litellm/issues/24375
+
+    When a PDF is sent to Gemini, the API returns a DOCUMENT modality in
+    promptTokensDetails. These tokens should be mapped to text_tokens since
+    Gemini bills documents at the text token rate. Without this fix, DOCUMENT
+    tokens are silently dropped, causing cost undercounting.
+
+    Real Gemini response for a PDF input:
+        promptTokensDetails: [TEXT: 8, DOCUMENT: 774]
+        candidatesTokensDetails: [TEXT: 4]
+        thoughtsTokenCount: 92
+    """
+    v = VertexGeminiConfig()
+
+    usage_metadata_dict = {
+        "promptTokenCount": 782,
+        "candidatesTokenCount": 4,
+        "totalTokenCount": 878,
+        "promptTokensDetails": [
+            {"modality": "TEXT", "tokenCount": 8},
+            {"modality": "DOCUMENT", "tokenCount": 774},
+        ],
+        "candidatesTokensDetails": [
+            {"modality": "TEXT", "tokenCount": 4},
+        ],
+        "thoughtsTokenCount": 92,
+    }
+
+    completion_response = {"usageMetadata": usage_metadata_dict}
+    result = v._calculate_usage(completion_response=completion_response)
+
+    # Verify basic token counts
+    assert result.prompt_tokens == 782
+    assert result.completion_tokens == 96  # 4 candidates + 92 thinking
+    assert result.total_tokens == 878
+
+    # DOCUMENT tokens should be included in text_tokens: 8 (TEXT) + 774 (DOCUMENT) = 782
+    assert result.prompt_tokens_details is not None
+    assert result.prompt_tokens_details.text_tokens == 782, \
+        "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)"
+
+    # Verify completion token details
+    assert result.completion_tokens_details is not None
+    assert result.completion_tokens_details.text_tokens == 4
+    assert result.completion_tokens_details.reasoning_tokens == 92
+
+
+def test_vertex_ai_usage_metadata_with_document_tokens_cached():
+    """Test that cached DOCUMENT tokens are correctly subtracted from prompt text tokens."""
+    v = VertexGeminiConfig()
+
+    usage_metadata_dict = {
+        "promptTokenCount": 782,
+        "candidatesTokenCount": 4,
+        "totalTokenCount": 878,
+        "cachedContentTokenCount": 400,
+        "promptTokensDetails": [
+            {"modality": "TEXT", "tokenCount": 8},
+            {"modality": "DOCUMENT", "tokenCount": 774},
+        ],
+        "cacheTokensDetails": [
+            {"modality": "DOCUMENT", "tokenCount": 400},
+        ],
+        "candidatesTokensDetails": [
+            {"modality": "TEXT", "tokenCount": 4},
+        ],
+        "thoughtsTokenCount": 92,
+    }
+
+    completion_response = {"usageMetadata": usage_metadata_dict}
+    result = v._calculate_usage(completion_response=completion_response)
+
+    # DOCUMENT cached tokens map to cached_text_tokens, so:
+    # text_tokens = (8 TEXT + 774 DOCUMENT) - 400 cached = 382
+    assert result.prompt_tokens_details.text_tokens == 382, \
+        "text_tokens should be (8 + 774) - 400 cached = 382"
+    assert result.prompt_tokens_details.cached_tokens == 400
+
+
 def test_async_streaming_uses_custom_client():
     """
     Test that user-specified async client is correctly passed to make_call