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
This commit is contained in:
Chesars 2026-03-23 10:39:23 -03:00
parent f5194b5ce3
commit 446456b855
2 changed files with 94 additions and 0 deletions

View file

@ -1721,6 +1721,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
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 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
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 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
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 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
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

View file

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