fix(vertex_ai): exclude Gemini Google Search grounding tokens from input token billing

[devin-ai-integration]

Relevant issues

Follow-up to #33533 (which surfaced Gemini toolUsePromptTokenCount) and the cost concern raised in #33198

Linear ticket

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Live proxy (litellm/proxy/proxy_cli.py), real gemini/gemini-2.5-flash calls costing real spend. Token counts differ per call because they're independent live requests

Before, base litellm_internal_staging at e59add11cd (the #33533 behavior): Google Search grounding tool-use tokens are folded into prompt_tokens, so they get billed at the input token rate on top of the separate web search fee

$ curl -s http://localhost:4000/v1/chat/completions -H "Authorization: Bearer sk-1234" \
    -d '{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"Who won the last 3 Nobel Prizes in Physics? Search the web and cite sources."}],"tools":[{"googleSearch":{}}]}'
{
  "prompt_tokens": 185,          # 19 text + 166 tool_use  <- overcharged
  "completion_tokens": 518,
  "total_tokens": 703,
  "prompt_tokens_details": {"text_tokens": 19, "tool_use_tokens": 166, "web_search_requests": 3}
}

After, this branch at 7cc6262d22, Google Search grounding: tool-use tokens stay visible in prompt_tokens_details.tool_use_tokens but are excluded from prompt_tokens

$ curl -s http://localhost:4000/v1/chat/completions -H "Authorization: Bearer sk-1234" \
    -d '{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"Who won the last 3 Nobel Prizes in Physics? Search the web and cite sources."}],"tools":[{"googleSearch":{}}]}'
{
  "prompt_tokens": 19,           # text only  <- correct
  "completion_tokens": 395,
  "total_tokens": 563,
  "prompt_tokens_details": {"text_tokens": 19, "tool_use_tokens": 149, "web_search_requests": 3}
}

After, this branch at 7cc6262d22, URL context (a non-search server-side tool): tool-use tokens are billed as input tokens, so they remain folded into prompt_tokens

$ curl -s http://localhost:4000/v1/chat/completions -H "Authorization: Bearer sk-1234" \
    -d '{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"Summarize the main heading of https://ai.google.dev/gemini-api/docs/pricing"}],"tools":[{"urlContext":{}}]}'
{
  "prompt_tokens": 14138,        # 23 text + 14115 tool_use  <- correct
  "completion_tokens": 346,
  "total_tokens": 14484,
  "prompt_tokens_details": {"text_tokens": 23, "tool_use_tokens": 14115}
}

Type

🐛 Bug Fix

Changes

Per Google's pricing (https://ai.google.dev/gemini-api/docs/pricing), retrieved context from Grounding with Google Search is not charged as input tokens; you pay a separate per-request / per-query search fee instead. Other server-side tools (URL context, File Search, code execution) do bill their tool-use tokens at the input token rate. toolUsePromptTokenCount is the same aggregate field for all of them, so #33533's unconditional prompt_tokens = promptTokenCount + toolUsePromptTokenCount overcharges search-grounded requests, double-counting the retrieved tokens once at the input rate and again via the flat search fee

This keys the fold-in on whether the response actually performed a Google Search, detected via groundingMetadata.webSearchQueries (the same signal that drives web_search_requests and the search fee). URL context also emits groundingMetadata, but with groundingChunks and no webSearchQueries, so presence of groundingMetadata alone is not sufficient

billable_tool_use_prompt_tokens = (
    0
    if VertexGeminiConfig._response_has_search_grounding(completion_response)
    else (tool_use_prompt_tokens or 0)
)
usage = Usage(
    prompt_tokens=usage_metadata.get("promptTokenCount", 0) + billable_tool_use_prompt_tokens,
    ...
)

prompt_tokens_details.tool_use_tokens is still populated in every case for observability, so the total_tokens gap that search grounding introduces stays explainable. The new _response_has_search_grounding helper reuses the existing _extract_candidate_metadata and _calculate_web_search_requests, so both the streaming and non-streaming usage paths share one decision

Tests in test_vertex_and_google_ai_studio_gemini.py cover search grounding (tool tokens excluded from prompt_tokens, still surfaced), URL context with groundingChunks but no webSearchQueries (tool tokens included), and the detection helper itself

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Link to Devin session: https://app.devin.ai/sessions/25ff7523b62f4304a1dde7132834ad6a
Requested by: @krrish-berri-2

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
0 out of 2 committers have signed the CLA.

❌ krrish-berri
❌ devin-ai-integration[bot]
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes a billing overcharge for Gemini requests that use Grounding with Google Search: previously, toolUsePromptTokenCount was unconditionally added to prompt_tokens, doubling the cost for search-retrieved tokens on top of the separate per-query search fee. The fix introduces _response_has_search_grounding which inspects groundingMetadata.webSearchQueries and excludes tool-use tokens from prompt_tokens only for search-grounded responses, leaving URL-context and other server-side tools unaffected.

• Detection logic (_response_has_search_grounding): correctly distinguishes Google Search (has webSearchQueries) from URL context (has groundingChunks but no webSearchQueries), and is covered by unit tests.
• Non-streaming path: correctly excludes tool-use tokens from prompt_tokens while still surfacing them on prompt_tokens_details.tool_use_tokens; live screenshots in the PR description confirm the fix.
• Streaming path: _calculate_usage is called with the per-chunk payload, and _response_has_search_grounding returns False when candidates is absent — creating a fragile dependency on Google's streaming format where usage and candidates arrive in the same chunk.
• total_tokens invariant: for search-grounded responses total_tokens now exceeds prompt_tokens + completion_tokens by tool_use_tokens; this is intentional but undocumented in the Usage construction and could mislead downstream cost accounting.

Confidence Score: 4/5

The non-streaming billing fix is correct and well-tested; the streaming path works under the current Gemini API format but contains a fragile dependency that could silently over-bill if the API delivers usage and candidates in separate events.

The core logic is sound and clearly fixes the overcharging bug for non-streaming requests. The _response_has_search_grounding helper correctly distinguishes search grounding from URL context, and three focused unit tests cover the key cases. Two design-level concerns remain: the streaming billing path derives search-grounding from the per-chunk payload rather than from the already-computed grounding_metadata list, making it implicitly dependent on how Google packs SSE events; and total_tokens diverges from prompt_tokens + completion_tokens for search-grounded calls without documentation at the construction site.

vertex_and_google_ai_studio_gemini.py around _apply_stream_usage_metadata and the Usage constructor call in _calculate_usage.

Important Files Changed

Filename	Overview
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py	Adds _response_has_search_grounding and uses it to exclude Google Search tool-use tokens from prompt_tokens; streaming path relies on an implicit assumption that usageMetadata and candidates (with grounding metadata) arrive in the same chunk, and total_tokens now exceeds prompt_tokens + completion_tokens for search-grounded responses.
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py	Adds three new unit tests covering search-grounding detection, search-grounding token exclusion, and URL-context token inclusion; all use mocked data (no network calls). No streaming-path test for the new billing logic.

Comments Outside Diff (1)

1. litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py, line 3247-3253 (link)

   Streaming path may silently over-bill search-grounding calls

   _calculate_usage is called with the raw streaming processed_chunk, and _response_has_search_grounding immediately returns False if "candidates" is absent from that chunk. In practice the Gemini API sends usageMetadata in the same SSE event as the final candidates batch (which contains groundingMetadata), so this usually works — but the existing grounding_metadata list, already extracted from the current chunk's candidates by the caller just a few lines above, is not forwarded into _calculate_usage. If the API ever delivers a usage-only final event (no candidates), the search-grounding check fails silently and toolUsePromptTokenCount is folded into prompt_tokens, reproducing the original over-billing bug for streaming consumers. The safer approach would be to compute billable_tool_use_prompt_tokens inside _apply_stream_usage_metadata using the already-available grounding_metadata argument — e.g., has_search = bool(_calculate_web_search_requests(grounding_metadata)) — instead of re-deriving it from the chunk.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "fix(vertex_ai): exclude Google Search gr..." | Re-trigger Greptile}

[greptile-apps]

total_tokens invariant broken for search-grounding responses

After this change, for search-grounding responses total_tokens (taken verbatim from totalTokenCount) will always be greater than prompt_tokens + completion_tokens by exactly the excluded toolUsePromptTokenCount. The test test_vertex_ai_search_grounding_tool_use_tokens_excluded_from_prompt_tokens explicitly asserts this gap. Any downstream code (cost calculators, budget checks, token accounting) that assumes total_tokens == prompt_tokens + completion_tokens will silently produce wrong results for search-grounded requests. Consider whether total_tokens should be adjusted to promptTokenCount + billable_tool_use_prompt_tokens + candidatesTokenCount + thoughtsTokenCount, or at minimum add a comment explaining the deliberate discrepancy in the Usage constructor call.

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_gemini_tool_use_input_cost (881a28f) with litellm_internal_staging (f759c75)}