From 446456b855e6cea909a0abe2082a7b50be5adc84 Mon Sep 17 00:00:00 2001 From: Chesars Date: Mon, 23 Mar 2026 10:39:23 -0300 Subject: [PATCH] fix(gemini): include DOCUMENT modality tokens in cost calculation Gemini API returns a DOCUMENT modality in promptTokensDetails for PDF inputs, but the token parser only handled TEXT, IMAGE, AUDIO, and VIDEO. DOCUMENT tokens were silently dropped, causing cost to be undercounted by up to 99% for PDF-heavy requests. Map DOCUMENT tokens to text_tokens since Gemini bills documents at the text token rate. Applied to all four modality parser loops: promptTokensDetails, cacheTokensDetails, responseTokensDetails, and candidatesTokensDetails. Fixes #24375 --- .../vertex_and_google_ai_studio_gemini.py | 12 +++ ...test_vertex_and_google_ai_studio_gemini.py | 82 +++++++++++++++++++ 2 files changed, 94 insertions(+) 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 36f51c5b2f59..b0e42d0a690d 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 @@ -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 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 3102a6959618..5c53b5831c18 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 @@ -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