fix(cost): bill Gemini audio/image/video prompt tokens at the resolved input rate when no modality rate exists (#47373)

* fix(cost): bill Gemini audio/image/video prompt tokens at the resolved input rate when no modality rate exists

Co-authored-by: shivam <shivam@berri.ai>

* style(cost): ruff format the modality cost calls

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: shivam <shivam@berri.ai>

Backport of #47373 to stable/1.106.x
Cherry-picked from 09e9193514 (main)

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -1086,6 +1086,25 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
     )
 
 
+def _modality_token_cost(
+    model_info: ModelInfo,
+    modality_cost_key: str,
+    tokens: int,
+    prompt_base_cost: float,
+    service_tier: str | None,
+) -> float:
+    """
+    Audio, image and video prompt tokens bill at their modality rate when the model has one.
+    Without one they are ordinary prompt tokens, so they bill at `prompt_base_cost`, which
+    `_get_token_base_cost` has already resolved for the service tier and long-context tier.
+    """
+    tiered_key: Final = get_service_tier_cost_key(modality_cost_key, service_tier)
+    cost_key: Final = tiered_key if model_info.get(tiered_key) is not None else modality_cost_key
+    if model_info.get(cost_key) is None:
+        return float(tokens) * prompt_base_cost
+    return calculate_cost_component(model_info, cost_key, tokens)
+
+
 def _calculate_input_cost(
     prompt_tokens_details: PromptTokensDetailsResult,
     model_info: ModelInfo,
@@ -1116,28 +1135,37 @@ def _calculate_input_cost(
     if prompt_tokens_details["audio_tokens"] and not (
         prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
     ):
-        audio_cost_key: Final = get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
-        prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
+        prompt_cost += _modality_token_cost(
+            model_info,
+            "input_cost_per_audio_token",
+            prompt_tokens_details["audio_tokens"],
+            prompt_base_cost,
+            service_tier,
+        )
 
     ### IMAGE TOKEN COST
     if prompt_tokens_details["image_tokens"] and not (
         prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
     ):
-        # For image token costs:
-        # First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
-        image_token_cost_key = "input_cost_per_image_token"
-        if model_info.get(image_token_cost_key) is None:
-            image_token_cost_key = "input_cost_per_token"
-        prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
+        prompt_cost += _modality_token_cost(
+            model_info,
+            "input_cost_per_image_token",
+            prompt_tokens_details["image_tokens"],
+            prompt_base_cost,
+            service_tier,
+        )
 
     ### VIDEO TOKEN COST
     if prompt_tokens_details["video_tokens"] and not (
         prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
     ):
-        video_token_cost_key = "input_cost_per_video_token"
-        if model_info.get(video_token_cost_key) is None:
-            video_token_cost_key = "input_cost_per_token"
-        prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"])
+        prompt_cost += _modality_token_cost(
+            model_info,
+            "input_cost_per_video_token",
+            prompt_tokens_details["video_tokens"],
+            prompt_base_cost,
+            service_tier,
+        )
 
     ### CACHE WRITING COST - Now uses tiered pricing
     if (

tests/unit/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py

@@ -450,6 +450,66 @@ def test_video_tokens_fallback_to_base_cost():
     assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12)
 
 
+def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_input_rate():
+    """Prompt audio/video tokens with no modality rate are prompt tokens, so a >200k prompt bills them above-200k."""
+    from unittest.mock import patch
+
+    mock_model_info = {
+        "input_cost_per_token": 2e-6,
+        "input_cost_per_token_above_200k_tokens": 4e-6,
+        "output_cost_per_token": 1.2e-5,
+        "output_cost_per_token_above_200k_tokens": 1.8e-5,
+    }
+
+    usage = Usage(
+        prompt_tokens=816032,
+        completion_tokens=4559,
+        total_tokens=820591,
+        prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=9033, audio_tokens=14999, video_tokens=792000),
+        completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=4559),
+    )
+
+    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="vertex_ai"
+        )
+
+    assert round(prompt_cost, 9) == round(816032 * 4e-6, 9)
+    assert round(completion_cost, 9) == round(4559 * 1.8e-5, 9)
+
+
+def test_prompt_modality_rates_still_win_over_tiered_input_rate():
+    """An explicit input_cost_per_audio_token / input_cost_per_video_token beats the fallback, even above 200k."""
+    from unittest.mock import patch
+
+    mock_model_info = {
+        "input_cost_per_token": 2e-6,
+        "input_cost_per_token_above_200k_tokens": 4e-6,
+        "input_cost_per_audio_token": 3e-6,
+        "input_cost_per_video_token": 1e-6,
+        "output_cost_per_token": 1.2e-5,
+    }
+
+    usage = Usage(
+        prompt_tokens=300000,
+        completion_tokens=10,
+        total_tokens=300010,
+        prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100000, audio_tokens=100000, video_tokens=100000),
+        completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10),
+    )
+
+    with patch(
+        "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
+        return_value=mock_model_info,
+    ):
+        prompt_cost, _ = generic_cost_per_token(model="test-model", usage=usage, custom_llm_provider="vertex_ai")
+
+    assert round(prompt_cost, 9) == round(100000 * 4e-6 + 100000 * 3e-6 + 100000 * 1e-6, 9)
+
+
 def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map):
     # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing
     model = "gemini-2.5-pro"