From 8b9ffdd93f62721c8f3871265d32ae24a6396b2d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 12:47:36 +0530 Subject: [PATCH] feat(vertex-ai): add VIDEO modality support in token usage tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Parse VIDEO modality in promptTokensDetails → prompt_tokens_details.video_tokens - Parse VIDEO modality in candidatesTokensDetails → completion_tokens_details.video_tokens - Parse VIDEO modality in cacheTokensDetails and subtract from prompt video tokens - Add video_tokens field to PromptTokensDetailsWrapper and CompletionTokensDetailsWrapper - Fix implicit caching text-token fallback to not fire when cacheTokensDetails is present - Add 4 unit tests covering: prompt video tokens, response video tokens, auto-calculated text fallback with video, and explicit video cache subtraction Co-Authored-By: Claude Sonnet 4.6 --- .../vertex_and_google_ai_studio_gemini.py | 16 +- ...odel_prices_and_context_window_backup.json | 2 +- litellm/types/utils.py | 6 + ...test_vertex_and_google_ai_studio_gemini.py | 150 ++++++++++++++++++ 4 files changed, 172 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 7bcefc1dd8..6fafbd3eda 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -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( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f52288ea72..cbd64a178b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16289,7 +16289,7 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", + "litellm_provider": "gemini", "max_audio_length_hours": 8.4, "max_audio_per_prompt": 1, "supports_reasoning": false, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8b9359876e..503817054a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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.""" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 6047da66b6..3a12c8a0bf 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -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 +