mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
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:
parent
f5194b5ce3
commit
446456b855
2 changed files with 94 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue