fix(gemini): resolve image token undercounting in usage metadata

[gustipardo]

This PR fixes the issue where image_tokens were being overwritten instead of accumulated in the usage metadata for Gemini/Vertex AI models.

Changes:

• Implemented token accumulation (+=) instead of overwriting.
• Added support for both 'tokenCount' and 'token_count' keys.
• Normalized modality strings to uppercase for consistency.

Verified with a standalone reproduction script (190 tokens counted vs 100 previously).

Fixes #22082

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Error		Mar 3, 2026 3:23pm

[CLAassistant]

All committers have signed the CLA.

[greptile-apps]

Greptile Summary

This PR fixes image token undercounting in Gemini/Vertex AI usage metadata by changing token assignment from overwriting (=) to accumulation (+=). When the API returns multiple entries for the same modality (e.g., two IMAGE entries in promptTokensDetails), tokens are now correctly summed. The fix also adds defensive handling for both tokenCount and token_count API key formats, and normalizes modality strings to uppercase for consistency.

• Token accumulation fix: All four token detail loops in VertexGeminiConfig._calculate_usage (responseTokensDetails, candidatesTokensDetails, promptTokensDetails, cacheTokensDetails) and the GoogleImageGenConfig._transform_image_usage method now accumulate instead of overwrite
• Dual-key support: A _get_token_count helper tries tokenCount first, then falls back to token_count for API compatibility
• Test coverage: Two new tests verify the accumulation behavior — one for the image generation path and one for the vertex _calculate_usage path, both using mock data only

Confidence Score: 4/5

• This PR is safe to merge — it fixes a clear bug with a well-scoped change and includes regression tests.
• Score of 4 reflects a focused, correct bug fix with good test coverage. The changes are straightforward accumulation fixes that don't alter control flow. Both the image generation and vertex calculation paths are covered by tests. Slight deduction because the similar code in vertex_gemini_transformation.py was not updated with the same defensive improvements (token_count fallback, .upper() normalization), though that's existing code not broken by this PR.
• No files require special attention. The main logic file litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py has the most changes but they are consistent and well-tested.

Important Files Changed

Filename	Overview
litellm/llms/gemini/image_generation/transformation.py	Changed token assignment from = to += for accumulation, added token_count key fallback and .upper() normalization. Clean and correct fix.
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py	Added _get_token_count helper for dual-key lookup, updated all four token detail loops to accumulate instead of overwrite. Logic is correct and consistent.
tests/llm_translation/test_gemini_image_usage.py	Added regression test for image token accumulation in GoogleImageGenConfig. No network calls, uses local model cost map. Covers the core fix well.
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py	Added test for VertexGeminiConfig._calculate_usage with duplicate modality entries. Tests both tokenCount and token_count keys, and verifies cache subtraction logic.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[API Response with usageMetadata] --> B{Parse token details loops}
    B --> C[responseTokensDetails]
    B --> D[candidatesTokensDetails]
    B --> E[promptTokensDetails]
    B --> F[cacheTokensDetails]
    
    C --> G[_get_token_count: try tokenCount then token_count]
    D --> G
    E --> G
    F --> G
    
    G --> H[Normalize modality to UPPERCASE]
    H --> I{Accumulate tokens by modality}
    I --> |TEXT| J["field = (field or 0) + token_count"]
    I --> |IMAGE| J
    I --> |AUDIO| J
    I --> |VIDEO| J
    
    E --> K[prompt_*_tokens accumulated]
    F --> L[cached_*_tokens accumulated]
    K --> M[Subtract cached from prompt per modality]
    L --> M
    M --> N[Final Usage object]
    C --> N
    D --> N

Loading

_{Last reviewed commit: de63dd8}

[greptile-apps]

_{5 files reviewed, 5 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Missing test for vertex _calculate_usage accumulation

This test covers the GoogleImageGenConfig._transform_image_usage path (image generation), but the PR also changed token accumulation logic in VertexGeminiConfig._calculate_usage across four loops (responseTokensDetails, candidatesTokensDetails, promptTokensDetails, cacheTokensDetails). None of those accumulation changes are covered by a test with duplicate modality entries.

Consider adding a test in tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py that passes promptTokensDetails (or candidatesTokensDetails) with multiple entries of the same modality to verify the accumulation works there too.