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

[Chesars]

Relevant issues

Fixes #24375

Pre-Submission checklist

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Type

🐛 Bug Fix

Changes

When a PDF is sent to a Gemini model, the API returns a DOCUMENT modality in promptTokensDetails (e.g., {modality: "DOCUMENT", tokenCount: 774}). The Gemini response parser only handled TEXT, IMAGE, AUDIO, and VIDEO modalities — DOCUMENT was silently dropped.

This caused prompt_tokens_details.text_tokens to report only the text portion (e.g., 9 of 783 tokens), and cost calculation used that incomplete breakdown, undercounting spend by up to 99% for PDF-heavy requests.

Fix

Map DOCUMENT modality tokens to text_tokens (Gemini bills documents at the text token rate). Applied to all four modality parser loops in vertex_and_google_ai_studio_gemini.py:

• promptTokensDetails → prompt_text_tokens
• cacheTokensDetails → cached_text_tokens
• responseTokensDetails → response_tokens_details.text_tokens
• candidatesTokensDetails → response_tokens_details.text_tokens

Tests added

• test_vertex_ai_usage_metadata_with_document_tokens_in_prompt — verifies DOCUMENT tokens are accumulated into text_tokens
• test_vertex_ai_usage_metadata_with_document_tokens_cached — verifies cached DOCUMENT tokens are correctly subtracted

[vercel]

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

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Mar 23, 2026 1:45pm

[greptile-apps]

Greptile Summary

This PR fixes a cost undercounting bug for Gemini models when PDFs are included in prompts. The Gemini API returns a DOCUMENT modality in token detail fields (e.g., promptTokensDetails, cacheTokensDetails), but the parser silently dropped it, causing only the small text-token portion to be counted while the bulk of PDF tokens went unaccounted — potentially undercounting costs by up to 99%.

Key changes:

• Maps DOCUMENT modality tokens to text_tokens in all four token-detail parsing loops in _calculate_usage (responseTokensDetails, candidatesTokensDetails, promptTokensDetails, cacheTokensDetails), consistent with Gemini's billing policy of charging document tokens at the text rate.
• Adds two focused unit tests covering the non-cached and cached DOCUMENT token paths; both are purely mock-based with no network calls, satisfying the test policy for this directory.
• The fix is minimal, symmetrical, and well-contained within the llms/ directory as expected per the provider-specific code policy.

Confidence Score: 5/5

• This PR is safe to merge — it is a narrow, additive bug fix with no risk of regressions to existing modality handling.
• The change is strictly additive (new elif branches only), touches a single well-isolated function, and is directly verified by two new unit tests that exercise both the non-cached and cached DOCUMENT token paths. No existing logic is altered or removed, and the mapping of DOCUMENT → text_tokens is consistent with all other loops in the same function as well as Gemini's published billing semantics.
• No files require special attention.

Important Files Changed

Filename	Overview
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py	Adds DOCUMENT modality handling to all four token-detail parsing loops (responseTokensDetails, candidatesTokensDetails, promptTokensDetails, cacheTokensDetails), mapping DOCUMENT tokens to text_tokens to match Gemini's billing semantics. Fix is minimal, consistent, and correctly placed.
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py	Adds two new unit tests: one verifying DOCUMENT tokens accumulate into text_tokens from promptTokensDetails, and one verifying cached DOCUMENT tokens are correctly subtracted. Both tests are purely mock-based with no network calls.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Gemini API Response\nusageMetadata] --> B{responseTokensDetails\npresent?}
    B -- Yes --> C[Loop over details]
    C --> D{modality?}
    D -- TEXT --> E[text_tokens += count]
    D -- AUDIO --> F[audio_tokens += count]
    D -- DOCUMENT\n✅ NEW --> E

    A --> G{candidatesTokensDetails\npresent?}
    G -- Yes --> H[Loop over details]
    H --> I{modality?}
    I -- TEXT --> J[text_tokens += count]
    I -- AUDIO --> K[audio_tokens += count]
    I -- IMAGE --> L[image_tokens += count]
    I -- VIDEO --> M[video_tokens += count]
    I -- DOCUMENT\n✅ NEW --> J

    A --> N{promptTokensDetails\npresent?}
    N -- Yes --> O[Loop over details]
    O --> P{modality?}
    P -- TEXT --> Q[prompt_text_tokens += count]
    P -- AUDIO --> R[prompt_audio_tokens += count]
    P -- IMAGE --> S[prompt_image_tokens += count]
    P -- VIDEO --> T[prompt_video_tokens += count]
    P -- DOCUMENT\n✅ NEW --> Q

    A --> U{cacheTokensDetails\npresent?}
    U -- Yes --> V[Loop over details]
    V --> W{modality?}
    W -- TEXT --> X[cached_text_tokens += count]
    W -- AUDIO --> Y[cached_audio_tokens += count]
    W -- IMAGE --> Z[cached_image_tokens += count]
    W -- VIDEO --> AA[cached_video_tokens += count]
    W -- DOCUMENT\n✅ NEW --> X

    Q --> AB[Subtract cached\nfrom prompt tokens]
    X --> AB
    AB --> AC[Build PromptTokensDetailsWrapper]
    J --> AD[Build CompletionTokensDetailsWrapper]
    AC --> AE[Usage object\nwith correct costs]
    AD --> AE

Loading

_{Reviews (1): Last reviewed commit: "fix(gemini): include DOCUMENT modality t..." | Re-trigger Greptile}

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing Chesars:fix/gemini-document-modality-cost (446456b) with main (c89496f)}