tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 04:21:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Gemini models: capture image_tokens and support cost_per_output_image_token in costs calculations (#16912)

Browse Source
This commit is contained in:
Dima-Mediator
2025-11-21 19:59:24 -08:00
committed by GitHub
parent 42c883d64f
commit a0d4d0b304
8 changed files with 257 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/llm_cost_calc/utils.py
+23
View File
@@ -408,6 +408,7 @@ class CompletionTokensDetailsResult(TypedDict):
audio_tokens: int
text_tokens: int
reasoning_tokens: int
image_tokens: int
def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult:
@@ -432,11 +433,19 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes
)
or 0
)
image_tokens = (
cast(
Optional[int],
getattr(usage.completion_tokens_details, "image_tokens", 0),
)
or 0
)
return CompletionTokensDetailsResult(
audio_tokens=audio_tokens,
text_tokens=text_tokens,
reasoning_tokens=reasoning_tokens,
image_tokens=image_tokens,
)
@@ -565,12 +574,14 @@ def generic_cost_per_token(
text_tokens = 0
audio_tokens = 0
reasoning_tokens = 0
image_tokens = 0
is_text_tokens_total = False
if usage.completion_tokens_details is not None:
completion_tokens_details = _parse_completion_tokens_details(usage)
audio_tokens = completion_tokens_details["audio_tokens"]
text_tokens = completion_tokens_details["text_tokens"]
reasoning_tokens = completion_tokens_details["reasoning_tokens"]
image_tokens = completion_tokens_details["image_tokens"]
if text_tokens == 0:
text_tokens = usage.completion_tokens
@@ -585,6 +596,9 @@ def generic_cost_per_token(
_output_cost_per_reasoning_token = _get_cost_per_unit(
model_info, "output_cost_per_reasoning_token", None
)
_output_cost_per_image_token = _get_cost_per_unit(
model_info, "output_cost_per_image_token", None
)
## AUDIO COST
if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0:
@@ -604,6 +618,15 @@ def generic_cost_per_token(
)
completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token
## IMAGE COST
if not is_text_tokens_total and image_tokens and image_tokens > 0:
_output_cost_per_image_token = (
_output_cost_per_image_token
if _output_cost_per_image_token is not None
else completion_base_cost
)
completion_cost += float(image_tokens) * _output_cost_per_image_token
return prompt_cost, completion_cost
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+28
View File
@@ -1379,6 +1379,30 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
response_tokens_details.audio_tokens = detail.get("tokenCount", 0)
#########################################################
## CANDIDATES TOKEN DETAILS (e.g., for image generation models) ##
if "candidatesTokensDetails" in usage_metadata:
if response_tokens_details is None:
response_tokens_details = CompletionTokensDetailsWrapper()
for detail in usage_metadata["candidatesTokensDetails"]:
modality = detail.get("modality")
token_count = detail.get("tokenCount", 0)
if modality == "TEXT":
response_tokens_details.text_tokens = token_count
elif modality == "AUDIO":
response_tokens_details.audio_tokens = token_count
elif modality == "IMAGE":
response_tokens_details.image_tokens = token_count
# Calculate text_tokens if not explicitly provided in candidatesTokensDetails
# candidatesTokenCount includes all modalities, so: text = total - (image + audio)
if response_tokens_details.text_tokens is None:
candidates_token_count = usage_metadata.get("candidatesTokenCount", 0)
image_tokens = response_tokens_details.image_tokens or 0
audio_tokens_candidate = response_tokens_details.audio_tokens or 0
calculated_text_tokens = candidates_token_count - image_tokens - audio_tokens_candidate
response_tokens_details.text_tokens = calculated_text_tokens
#########################################################
if "promptTokensDetails" in usage_metadata:
for detail in usage_metadata["promptTokensDetails"]:
if detail["modality"] == "AUDIO":
@@ -1387,6 +1411,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
text_tokens = detail.get("tokenCount", 0)
if "thoughtsTokenCount" in usage_metadata:
reasoning_tokens = usage_metadata["thoughtsTokenCount"]
# Also add reasoning tokens to response_tokens_details
if response_tokens_details is None:
response_tokens_details = CompletionTokensDetailsWrapper()
response_tokens_details.reasoning_tokens = reasoning_tokens
## adjust 'text_tokens' to subtract cached tokens
if (
litellm/types/utils.py
+30 -8
View File
@@ -162,6 +162,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
float
] # only for vertex ai models
output_cost_per_image: Optional[float]
output_cost_per_image_token: Optional[float]
output_vector_size: Optional[int]
output_cost_per_reasoning_token: Optional[float]
output_cost_per_video_per_second: Optional[float] # only for vertex ai models
@@ -945,6 +946,9 @@ class CompletionTokensDetailsWrapper(
text_tokens: Optional[int] = None
"""Text tokens generated by the model."""
image_tokens: Optional[int] = None
"""Image tokens generated by the model."""
class CacheCreationTokenDetails(BaseModel):
ephemeral_5m_input_tokens: Optional[int] = None
@@ -1033,15 +1037,8 @@ class Usage(CompletionUsage):
):
# handle reasoning_tokens
_completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None
if reasoning_tokens:
text_tokens = (
completion_tokens - reasoning_tokens if completion_tokens else None
)
completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens, text_tokens=text_tokens
)
# Ensure completion_tokens_details is properly handled
# First, handle existing completion_tokens_details
if completion_tokens_details:
if isinstance(completion_tokens_details, dict):
_completion_tokens_details = CompletionTokensDetailsWrapper(
@@ -1050,6 +1047,30 @@ class Usage(CompletionUsage):
elif isinstance(completion_tokens_details, CompletionTokensDetails):
_completion_tokens_details = completion_tokens_details
# Handle reasoning_tokens and auto-calculate text_tokens if needed
if reasoning_tokens:
# Ensure we have a details object to work with
if _completion_tokens_details is None:
_completion_tokens_details = CompletionTokensDetailsWrapper()
# Set reasoning_tokens if not already set by provider
if _completion_tokens_details.reasoning_tokens is None:
_completion_tokens_details.reasoning_tokens = reasoning_tokens
# Auto-calculate text_tokens only if provider didn't set it explicitly
# Formula: text_tokens = completion_tokens - reasoning_tokens - image_tokens - audio_tokens
if _completion_tokens_details.text_tokens is None and completion_tokens is not None:
calculated_text_tokens = completion_tokens - reasoning_tokens
# Subtract other modality tokens if present
if _completion_tokens_details.image_tokens:
calculated_text_tokens -= _completion_tokens_details.image_tokens
if _completion_tokens_details.audio_tokens:
calculated_text_tokens -= _completion_tokens_details.audio_tokens
# Prevent negative token counts from inconsistent data
_completion_tokens_details.text_tokens = max(0, calculated_text_tokens)
# handle prompt_tokens_details
_prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
@@ -2370,6 +2391,7 @@ class CustomPricingLiteLLMParams(BaseModel):
output_cost_per_token_above_200k_tokens: Optional[float] = None
output_cost_per_character_above_128k_tokens: Optional[float] = None
output_cost_per_image: Optional[float] = None
output_cost_per_image_token: Optional[float] = None
output_cost_per_reasoning_token: Optional[float] = None
output_cost_per_video_per_second: Optional[float] = None
output_cost_per_audio_per_second: Optional[float] = None
litellm/utils.py
+1
View File
@@ -5025,6 +5025,7 @@ def _get_model_info_helper( # noqa: PLR0915
"output_cost_per_video_per_second", None
),
output_cost_per_image=_model_info.get("output_cost_per_image", None),
output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None),
output_vector_size=_model_info.get("output_vector_size", None),
citation_cost_per_token=_model_info.get(
"citation_cost_per_token", None
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+92
View File
@@ -115,6 +115,98 @@ def test_reasoning_tokens_gemini():
)
def test_image_tokens_with_custom_pricing():
"""Test that image_tokens in completion are properly costed with output_cost_per_image_token."""
from unittest.mock import patch
# Mock model info with image token pricing
mock_model_info = {
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
"output_cost_per_image_token": 5e-6, # Custom pricing for image tokens in output
}
usage = Usage(
completion_tokens=1720, # text_tokens (600) + image_tokens (1120)
prompt_tokens=14,
total_tokens=1734,
completion_tokens_details=CompletionTokensDetailsWrapper(
accepted_prediction_tokens=None,
audio_tokens=None,
reasoning_tokens=0,
rejected_prediction_tokens=None,
text_tokens=600,
image_tokens=1120,
),
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None, cached_tokens=None, text_tokens=14, image_tokens=None
),
)
with patch(
"litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
return_value=mock_model_info,
):
prompt_cost, completion_cost = generic_cost_per_token(
model="test-model", usage=usage, custom_llm_provider="gemini"
)
# Expected costs:
# Prompt: 14 * 1e-6
# Completion: (600 * 2e-6) + (1120 * 5e-6)
expected_prompt_cost = 14 * 1e-6
expected_completion_cost = (600 * 2e-6) + (1120 * 5e-6)
assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
assert round(completion_cost, 12) == round(expected_completion_cost, 12)
def test_image_tokens_fallback_to_base_cost():
"""Test that image_tokens fall back to base cost when output_cost_per_image_token is not set."""
from unittest.mock import patch
# Mock model info without image token pricing
mock_model_info = {
"input_cost_per_token": 1e-6,
"output_cost_per_token": 2e-6,
# No output_cost_per_image_token defined
}
usage = Usage(
completion_tokens=1720,
prompt_tokens=14,
total_tokens=1734,
completion_tokens_details=CompletionTokensDetailsWrapper(
accepted_prediction_tokens=None,
audio_tokens=None,
reasoning_tokens=0,
rejected_prediction_tokens=None,
text_tokens=600,
image_tokens=1120,
),
prompt_tokens_details=PromptTokensDetailsWrapper(
audio_tokens=None, cached_tokens=None, text_tokens=14, image_tokens=None
),
)
with patch(
"litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
return_value=mock_model_info,
):
prompt_cost, completion_cost = generic_cost_per_token(
model="test-model", usage=usage, custom_llm_provider="gemini"
)
# Expected costs:
# Prompt: 14 * 1e-6
# Completion: (600 * 2e-6) + (1120 * 2e-6) # image_tokens use base cost
expected_prompt_cost = 14 * 1e-6
expected_completion_cost = (600 * 2e-6) + (1120 * 2e-6)
assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
assert round(completion_cost, 12) == round(expected_completion_cost, 12)
def test_generic_cost_per_token_above_200k_tokens():
model = "gemini-2.5-pro-exp-03-25"
custom_llm_provider = "vertex_ai"
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+80
View File
@@ -477,6 +477,86 @@ def test_vertex_ai_usage_metadata_response_token_count():
assert result.completion_tokens_details.text_tokens == 74
def test_vertex_ai_usage_metadata_with_image_tokens():
"""Test candidatesTokensDetails with IMAGE modality (e.g., Imagen models)
This test simulates the case where candidatesTokenCount is EXCLUSIVE of thoughtsTokenCount.
Gemini API returns: totalTokenCount = promptTokenCount + candidatesTokenCount + thoughtsTokenCount
"""
v = VertexGeminiConfig()
usage_metadata = {
"promptTokenCount": 14,
"candidatesTokenCount": 1442, # Does NOT include thoughtsTokenCount
"totalTokenCount": 1614, # 14 + 1442 + 158
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 14}],
"candidatesTokensDetails": [
{"modality": "IMAGE", "tokenCount": 1120},
{"modality": "TEXT", "tokenCount": 322} # 1442 - 1120 = 322
],
"thoughtsTokenCount": 158
}
usage_metadata = UsageMetadata(**usage_metadata)
result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata})
print("result", result)
# Verify basic token counts
assert result.prompt_tokens == 14
# completion_tokens = candidatesTokenCount + thoughtsTokenCount (when exclusive)
assert result.completion_tokens == 1600 # 1442 + 158
assert result.total_tokens == 1614
# Verify detailed token breakdown
assert result.completion_tokens_details.image_tokens == 1120
assert result.completion_tokens_details.text_tokens == 322
assert result.completion_tokens_details.reasoning_tokens == 158
# Verify the math: completion_tokens = image + text + reasoning
# 1600 = 1120 (image) + 322 (text) + 158 (reasoning)
assert (
result.completion_tokens_details.image_tokens
+ result.completion_tokens_details.text_tokens
+ result.completion_tokens_details.reasoning_tokens
== result.completion_tokens
)
def test_vertex_ai_usage_metadata_with_image_tokens_auto_calculated_text():
"""Test that text_tokens is auto-calculated when only IMAGE modality is provided
This test verifies the auto-calculation logic at line 1367-1372 in vertex_and_google_ai_studio_gemini.py
"""
v = VertexGeminiConfig()
usage_metadata = {
"promptTokenCount": 14,
"candidatesTokenCount": 1442,
"totalTokenCount": 1614, # 14 + 1442 + 158 (exclusive)
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 14}],
"candidatesTokensDetails": [
{"modality": "IMAGE", "tokenCount": 1120}
# TEXT modality omitted - should be auto-calculated
],
"thoughtsTokenCount": 158
}
usage_metadata = UsageMetadata(**usage_metadata)
result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata})
print("result", result)
# Verify basic token counts
assert result.prompt_tokens == 14
assert result.completion_tokens == 1600 # 1442 + 158
assert result.total_tokens == 1614
# Verify image_tokens is set
assert result.completion_tokens_details.image_tokens == 1120
# Verify text_tokens is auto-calculated: candidatesTokenCount - image_tokens
# Note: reasoning_tokens is NOT subtracted here because candidatesTokenCount is exclusive
# 1442 - 1120 = 322
expected_text_tokens = 1442 - 1120
assert result.completion_tokens_details.text_tokens == expected_text_tokens
assert result.completion_tokens_details.reasoning_tokens == 158
def test_vertex_ai_map_thinking_param_with_budget_tokens_0():
"""
If budget_tokens is 0, do not set includeThoughts to True
tests/test_litellm/test_utils.py
+1
View File
@@ -416,6 +416,7 @@ def validate_model_cost_values(model_data, exceptions=None):
"input_cost_per_request",
"input_cost_per_audio_token",
"output_cost_per_audio_token",
"output_cost_per_image_token",
"input_cost_per_audio_per_second",
"input_cost_per_video_per_second",
"input_cost_per_token_above_128k_tokens",
tests/test_litellm/types/test_types_utils.py
+2 -1
View File
@@ -119,7 +119,8 @@ def test_usage_completion_tokens_details_text_tokens():
'audio_tokens': None,
'reasoning_tokens': 65,
'rejected_prediction_tokens': None,
'text_tokens': 12
'text_tokens': 12,
'image_tokens': None
}
assert dump_result['completion_tokens_details'] == expected_completion_details