scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-13 03:59:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

fix(gemini): price tool-use prompt tokens as input (#46021)

Browse source 
* test(gemini): tool-use prompt tokens must be priced as input

Signed-off-by: Delweng <delweng@gmail.com>

* fix(gemini): price tool-use prompt tokens as input

toolUsePromptTokenCount was folded into usage.prompt_tokens but left out
of prompt_tokens_details.text_tokens, and generic_cost_per_token bills
input from text_tokens, so URL context / code execution / File Search
tool tokens were counted but never charged. Search-grounded responses
keep adding zero via the grounding-aware billable count, and
promptTokensDetails absent still leaves text_tokens None so billing
falls back to the full prompt_tokens

Signed-off-by: Delweng <delweng@gmail.com>

---------

Signed-off-by: Delweng <delweng@gmail.com>
This commit is contained in:
Delweng 2026-10-13 08:01:48 +08:00 • committed by GitHub
parent 4b7b8efd4c
commit b24b68e576
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 108 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -1958,20 +1958,23 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
tool_use_prompt_tokens: Final = usage_metadata.get("toolUsePromptTokenCount") or None
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cached_tokens,
audio_tokens=prompt_audio_tokens,
text_tokens=prompt_text_tokens,
image_tokens=prompt_image_tokens,
video_tokens=prompt_video_tokens,
tool_use_tokens=tool_use_prompt_tokens,
)
billable_tool_use_prompt_tokens: Final = (
0
if VertexGeminiConfig._response_has_search_grounding(completion_response)
else (tool_use_prompt_tokens or 0)
)
billable_text_tokens: Final = (
prompt_text_tokens if prompt_text_tokens is None else prompt_text_tokens + billable_tool_use_prompt_tokens
)
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=cached_tokens,
audio_tokens=prompt_audio_tokens,
text_tokens=billable_text_tokens,
image_tokens=prompt_image_tokens,
video_tokens=prompt_video_tokens,
tool_use_tokens=tool_use_prompt_tokens,
)
completion_tokens = response_tokens or completion_response["usageMetadata"].get("candidatesTokenCount", 0)
if not VertexGeminiConfig.is_candidate_token_count_inclusive(usage_metadata) and reasoning_tokens:

96
tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -812,6 +812,102 @@ def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens():
assert usage.total_tokens - usage.prompt_tokens - usage.completion_tokens == 0
def test_url_context_tool_use_tokens_are_priced_as_input():
"""
URL context tool-use tokens are billed at the input token rate, so the input cost must
cover every token counted in usage.prompt_tokens, not just the reported text modality.
See https://github.com/BerriAI/litellm/issues/45991
"""
v = VertexGeminiConfig()
completion_response = {
"candidates": [
{
"urlContextMetadata": {"urlMetadata": []},
"groundingMetadata": {
"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]
},
}
],
"usageMetadata": UsageMetadata(
promptTokenCount=23,
candidatesTokenCount=10,
toolUsePromptTokenCount=14115,
totalTokenCount=14148,
promptTokensDetails=[{"modality": "TEXT", "tokenCount": 23}],
toolUsePromptTokensDetails=[{"modality": "TEXT", "tokenCount": 14115}],
),
}
usage = v.calculate_usage(completion_response=completion_response)
assert usage.prompt_tokens == 23 + 14115
assert usage.prompt_tokens_details.text_tokens == usage.prompt_tokens
input_rate: Final = litellm.get_model_info("gemini-2.5-flash")["input_cost_per_token"]
prompt_cost, _ = litellm.cost_per_token(model="gemini/gemini-2.5-flash", usage_object=usage)
assert prompt_cost == usage.prompt_tokens * input_rate
def test_search_grounding_tool_use_tokens_stay_unpriced():
"""
Grounding with Google Search tokens are not billed at the input token rate (a separate
per-query search fee applies), so neither prompt_tokens nor the input cost may grow with
toolUsePromptTokenCount.
"""
v = VertexGeminiConfig()
completion_response = {
"candidates": [{"groundingMetadata": {"webSearchQueries": ["latest nobel physics"]}}],
"usageMetadata": UsageMetadata(
promptTokenCount=23,
candidatesTokenCount=10,
toolUsePromptTokenCount=14115,
totalTokenCount=14148,
promptTokensDetails=[{"modality": "TEXT", "tokenCount": 23}],
toolUsePromptTokensDetails=[{"modality": "TEXT", "tokenCount": 14115}],
),
}
usage = v.calculate_usage(completion_response=completion_response)
assert usage.prompt_tokens == 23
assert usage.prompt_tokens_details.text_tokens == 23
input_rate: Final = litellm.get_model_info("gemini-2.5-flash")["input_cost_per_token"]
prompt_cost, _ = litellm.cost_per_token(model="gemini/gemini-2.5-flash", usage_object=usage)
assert prompt_cost == 23 * input_rate
def test_missing_prompt_tokens_details_still_prices_the_full_prompt():
"""
Gemini omits promptTokensDetails on some tool-use responses, leaving text_tokens unset.
That None must be preserved so input billing falls back to the full prompt_tokens, and
must not be turned into a concrete count that prices only part of the prompt.
"""
v = VertexGeminiConfig()
completion_response = {
"candidates": [
{
"urlContextMetadata": {"urlMetadata": []},
"groundingMetadata": {
"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]
},
}
],
"usageMetadata": UsageMetadata(
promptTokenCount=23,
candidatesTokenCount=10,
toolUsePromptTokenCount=14115,
totalTokenCount=14148,
),
}
usage = v.calculate_usage(completion_response=completion_response)
assert usage.prompt_tokens == 23 + 14115
assert usage.prompt_tokens_details.text_tokens is None
input_rate: Final = litellm.get_model_info("gemini-2.5-flash")["input_cost_per_token"]
prompt_cost, _ = litellm.cost_per_token(model="gemini/gemini-2.5-flash", usage_object=usage)
assert prompt_cost == usage.prompt_tokens * input_rate
def test_streaming_chunk_includes_reasoning_tokens():
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,