ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #22550 from BerriAI/litellm_vertex-video-token-tracking

Browse Source 
feat(vertex-ai): add VIDEO modality support in token usage tracking
This commit is contained in:
Sameer Kankute 2026-03-02 18:51:19 +05:30 committed by GitHub
commit acf324279c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 171 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View File

@ -1590,6 +1590,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
prompt_audio_tokens: Optional[int] = None
prompt_image_tokens: Optional[int] = None
prompt_text_tokens: Optional[int] = None
prompt_video_tokens: Optional[int] = None
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
reasoning_tokens: Optional[int] = None
response_tokens: Optional[int] = None
@ -1624,9 +1625,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
response_tokens_details.audio_tokens = token_count
elif modality == "IMAGE":
response_tokens_details.image_tokens = token_count
elif modality == "VIDEO":
response_tokens_details.video_tokens = token_count
# Calculate text_tokens if not explicitly provided in candidatesTokensDetails
# candidatesTokenCount includes all modalities, so: text = total - (image + audio)
# candidatesTokenCount includes all modalities, so: text = total - (image + audio + video)
candidates_token_count = usage_metadata.get("candidatesTokenCount", 0)
if candidates_token_count > 0:
if response_tokens_details is None:
@ -1634,10 +1637,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if response_tokens_details.text_tokens is None:
completion_image_tokens = response_tokens_details.image_tokens or 0
completion_audio_tokens = response_tokens_details.audio_tokens or 0
completion_video_tokens = response_tokens_details.video_tokens or 0
calculated_text_tokens = (
candidates_token_count
- completion_image_tokens
- completion_audio_tokens
- completion_video_tokens
)
response_tokens_details.text_tokens = calculated_text_tokens
#########################################################
@ -1651,12 +1656,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
prompt_text_tokens = detail.get("tokenCount", 0)
elif detail["modality"] == "IMAGE":
prompt_image_tokens = detail.get("tokenCount", 0)
elif detail["modality"] == "VIDEO":
prompt_video_tokens = detail.get("tokenCount", 0)
## Parse cacheTokensDetails (breakdown of cached tokens by modality)
## When explicit caching is used, Gemini provides this field to show which modalities were cached
cached_text_tokens: Optional[int] = None
cached_audio_tokens: Optional[int] = None
cached_image_tokens: Optional[int] = None
cached_video_tokens: Optional[int] = None
if "cacheTokensDetails" in usage_metadata:
for detail in usage_metadata["cacheTokensDetails"]:
@ -1666,6 +1674,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
cached_text_tokens = detail.get("tokenCount", 0)
elif detail["modality"] == "IMAGE":
cached_image_tokens = detail.get("tokenCount", 0)
elif detail["modality"] == "VIDEO":
cached_video_tokens = detail.get("tokenCount", 0)
## Calculate non-cached tokens by subtracting cached from total (per modality)
## This is necessary because promptTokensDetails includes both cached and non-cached tokens
@ -1677,6 +1687,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
cached_tokens is not None
and prompt_text_tokens is not None
and cached_text_tokens is None
and "cacheTokensDetails" not in usage_metadata
):
# Implicit caching: only cachedContentTokenCount is provided (no cacheTokensDetails)
# Subtract from text tokens since implicit caching is primarily for text content
@ -1686,6 +1697,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens
if cached_image_tokens is not None and prompt_image_tokens is not None:
prompt_image_tokens = prompt_image_tokens - cached_image_tokens
if cached_video_tokens is not None and prompt_video_tokens is not None:
prompt_video_tokens = prompt_video_tokens - cached_video_tokens
if "thoughtsTokenCount" in usage_metadata:
reasoning_tokens = usage_metadata["thoughtsTokenCount"]
@ -1699,6 +1712,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
audio_tokens=prompt_audio_tokens,
text_tokens=prompt_text_tokens,
image_tokens=prompt_image_tokens,
video_tokens=prompt_video_tokens,
)
completion_tokens = response_tokens or completion_response["usageMetadata"].get(

6
litellm/types/utils.py
View File

@ -1383,6 +1383,9 @@ class CompletionTokensDetailsWrapper(
image_tokens: Optional[int] = None
"""Image tokens generated by the model."""
video_tokens: Optional[int] = None
"""Video tokens generated by the model."""
class CacheCreationTokenDetails(BaseModel):
ephemeral_5m_input_tokens: Optional[int] = None
@ -1398,6 +1401,9 @@ class PromptTokensDetailsWrapper(
image_tokens: Optional[int] = None
"""Image tokens sent to the model."""
video_tokens: Optional[int] = None
"""Video tokens sent to the model."""
web_search_requests: Optional[int] = None
"""Number of web search requests made by the tool call. Used for Anthropic to calculate web search cost."""

150
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View File

@ -3509,3 +3509,153 @@ def test_vertex_ai_web_search_options_in_map_openai_params():
assert optional_params["tools"][0]["googleSearch"] == {}, "googleSearch should be empty config"
assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation"
def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt():
"""Test promptTokensDetails with VIDEO modality for video inputs.
This test verifies that video tokens from promptTokensDetails are correctly
parsed and surfaced in prompt_tokens_details.video_tokens.
Based on a real Gemini response where a video file is sent as input:
promptTokensDetails: [VIDEO: 10240, TEXT: 9, AUDIO: 200]
candidatesTokensDetails: [TEXT: 79]
"""
v = VertexGeminiConfig()
usage_metadata_dict = {
"promptTokenCount": 10449,
"candidatesTokenCount": 79,
"totalTokenCount": 10528,
"trafficType": "ON_DEMAND",
"promptTokensDetails": [
{"modality": "VIDEO", "tokenCount": 10240},
{"modality": "TEXT", "tokenCount": 9},
{"modality": "AUDIO", "tokenCount": 200},
],
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": 79},
],
}
completion_response = {"usageMetadata": usage_metadata_dict}
result = v._calculate_usage(completion_response=completion_response)
# Verify basic token counts
assert result.prompt_tokens == 10449
assert result.completion_tokens == 79
assert result.total_tokens == 10528
# Verify prompt token details include video tokens
assert result.prompt_tokens_details is not None
assert result.prompt_tokens_details.video_tokens == 10240, \
"Prompt video tokens should be 10240"
assert result.prompt_tokens_details.text_tokens == 9, \
"Prompt text tokens should be 9"
assert result.prompt_tokens_details.audio_tokens == 200, \
"Prompt audio tokens should be 200"
# Verify completion token details
assert result.completion_tokens_details is not None
assert result.completion_tokens_details.text_tokens == 79, \
"Completion text tokens should be 79"
assert result.completion_tokens_details.video_tokens is None, \
"Completion video tokens should be None (text-only response)"
def test_vertex_ai_usage_metadata_with_video_tokens_in_candidates():
"""Test candidatesTokensDetails with VIDEO modality.
Verifies that video tokens in the response (candidatesTokensDetails) are
correctly parsed and reflected in completion_tokens_details.video_tokens,
and that text_tokens is auto-calculated by subtracting video tokens.
"""
v = VertexGeminiConfig()
usage_metadata_dict = {
"promptTokenCount": 10,
"candidatesTokenCount": 10330,
"totalTokenCount": 10340,
"promptTokensDetails": [
{"modality": "TEXT", "tokenCount": 10},
],
"candidatesTokensDetails": [
{"modality": "VIDEO", "tokenCount": 10240},
{"modality": "TEXT", "tokenCount": 90},
],
}
completion_response = {"usageMetadata": usage_metadata_dict}
result = v._calculate_usage(completion_response=completion_response)
assert result.completion_tokens == 10330
assert result.completion_tokens_details is not None
assert result.completion_tokens_details.video_tokens == 10240, \
"Completion video tokens should be 10240"
assert result.completion_tokens_details.text_tokens == 90, \
"Completion text tokens should be 90"
# Verify prompt side has no video tokens
assert result.prompt_tokens_details.video_tokens is None, \
"Prompt video tokens should be None (text-only input)"
def test_vertex_ai_usage_metadata_video_tokens_auto_calculated_text():
"""Test that text_tokens is auto-calculated correctly when VIDEO modality
is present in candidatesTokensDetails but TEXT is omitted.
text = candidatesTokenCount - video_tokens - image_tokens - audio_tokens
"""
v = VertexGeminiConfig()
usage_metadata_dict = {
"promptTokenCount": 10,
"candidatesTokenCount": 10330,
"totalTokenCount": 10340,
"candidatesTokensDetails": [
{"modality": "VIDEO", "tokenCount": 10240},
# TEXT intentionally omitted — should be auto-calculated
],
}
completion_response = {"usageMetadata": usage_metadata_dict}
result = v._calculate_usage(completion_response=completion_response)
assert result.completion_tokens_details.video_tokens == 10240
# text = 10330 - 10240 = 90
assert result.completion_tokens_details.text_tokens == 90, \
"text_tokens should be auto-calculated as candidatesTokenCount - video_tokens"
def test_vertex_ai_usage_metadata_video_tokens_with_caching():
"""Test that cached video tokens are correctly subtracted from prompt video tokens
when cacheTokensDetails includes VIDEO modality.
"""
v = VertexGeminiConfig()
usage_metadata_dict = {
"promptTokenCount": 10449,
"candidatesTokenCount": 79,
"totalTokenCount": 10528,
"cachedContentTokenCount": 5120,
"promptTokensDetails": [
{"modality": "VIDEO", "tokenCount": 10240},
{"modality": "TEXT", "tokenCount": 9},
{"modality": "AUDIO", "tokenCount": 200},
],
"cacheTokensDetails": [
{"modality": "VIDEO", "tokenCount": 5120},
],
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": 79},
],
}
completion_response = {"usageMetadata": usage_metadata_dict}
result = v._calculate_usage(completion_response=completion_response)
# video tokens should be reduced by cached amount: 10240 - 5120 = 5120
assert result.prompt_tokens_details.video_tokens == 5120, \
"Prompt video tokens should be 10240 - 5120 (cached) = 5120"
assert result.prompt_tokens_details.text_tokens == 9
assert result.prompt_tokens_details.audio_tokens == 200