Merge pull request #37724 from bisma-nawaz/fix-37647-staging

fix: map Gemini ON_DEMAND_FLEX traffic type to flex service tier

litellm/cost_calculator.py

@@ -591,6 +591,7 @@ def cost_per_token(
                 prompt_characters=prompt_characters,
                 completion_characters=completion_characters,
                 usage=usage_block,
+                service_tier=service_tier,
                 vertex_location=vertex_location,
             )
         elif cost_router == "cost_per_token":
@@ -845,9 +846,11 @@ def _get_response_model(completion_response: object) -> str | None:
 _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = {
     # ON_DEMAND_PRIORITY maps to "priority" — selects input_cost_per_token_priority, etc.
     "ON_DEMAND_PRIORITY": "priority",
-    # FLEX / BATCH maps to "flex" — selects input_cost_per_token_flex, etc.
+    # FLEX / BATCH / ON_DEMAND_FLEX maps to "flex" — selects input_cost_per_token_flex, etc.
+    # Vertex AI reports flex/shared-capacity traffic as ON_DEMAND_FLEX, not FLEX.
     "FLEX": "flex",
     "BATCH": "flex",
+    "ON_DEMAND_FLEX": "flex",
     # ON_DEMAND is standard pricing — no service_tier suffix applied
     "ON_DEMAND": None,
 }
@@ -862,9 +865,9 @@ def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None:
 
     trafficType values seen in practice
     ------------------------------------
-    ON_DEMAND          -> standard pricing  (service_tier = None)
-    ON_DEMAND_PRIORITY -> priority pricing  (service_tier = "priority")
-    FLEX / BATCH       -> batch/flex pricing (service_tier = "flex")
+    ON_DEMAND               -> standard pricing  (service_tier = None)
+    ON_DEMAND_PRIORITY      -> priority pricing  (service_tier = "priority")
+    FLEX / BATCH / ON_DEMAND_FLEX -> batch/flex pricing (service_tier = "flex")
     """
     if traffic_type is None:
         return None

litellm/llms/vertex_ai/cost_calculator.py

@@ -64,6 +64,7 @@ def cost_per_character(
     usage: Usage,
     prompt_characters: float | None = None,
     completion_characters: float | None = None,
+    service_tier: str | None = None,
     vertex_location: str | None = None,
 ) -> tuple[float, float]:
     """
@@ -74,6 +75,8 @@ def cost_per_character(
         - custom_llm_provider: str, "vertex_ai-*"
         - prompt_characters: float, the number of input characters
         - completion_characters: float, the number of output characters
+        - service_tier: optional tier derived from Gemini trafficType
+          ("priority" for ON_DEMAND_PRIORITY, "flex" for FLEX/batch).
         - vertex_location: the Vertex AI location serving the request; non-global
           locations apply the model's regional-endpoint uplift multiplier
 
@@ -92,6 +95,7 @@ def cost_per_character(
             model=model,
             custom_llm_provider=custom_llm_provider,
             usage=usage,
+            service_tier=service_tier,
         )
     else:
         try:
@@ -123,6 +127,7 @@ def cost_per_character(
                 model=model,
                 custom_llm_provider=custom_llm_provider,
                 usage=usage,
+                service_tier=service_tier,
             )
 
     ## CALCULATE OUTPUT COST
@@ -131,6 +136,7 @@ def cost_per_character(
             model=model,
             custom_llm_provider=custom_llm_provider,
             usage=usage,
+            service_tier=service_tier,
         )
     else:
         completion_tokens: Final = usage.completion_tokens
@@ -162,6 +168,7 @@ def cost_per_character(
                 model=model,
                 custom_llm_provider=custom_llm_provider,
                 usage=usage,
+                service_tier=service_tier,
             )
 
     vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)

tests/test_litellm/llms/gemini/test_cost_calculator.py

@@ -361,3 +361,33 @@ def test_gemini_image_generation_cost_no_web_search_when_absent(monkeypatch):
     )
 
     assert cost_zero == cost_none
+
+
+@pytest.mark.parametrize(
+    "traffic_type, expected_service_tier",
+    [
+        ("ON_DEMAND", None),
+        ("ON_DEMAND_PRIORITY", "priority"),
+        ("FLEX", "flex"),
+        ("BATCH", "flex"),
+        # Vertex AI reports flex/shared-capacity traffic as ON_DEMAND_FLEX.
+        ("ON_DEMAND_FLEX", "flex"),
+        # trafficType is matched case-insensitively.
+        ("on_demand_flex", "flex"),
+        (None, None),
+        ("SOMETHING_UNKNOWN", None),
+    ],
+)
+def test_map_traffic_type_to_service_tier(
+    traffic_type: str | None, expected_service_tier: str | None
+):
+    """
+    Gemini/Vertex usageMetadata.trafficType maps to the LiteLLM service_tier
+    that selects flex/priority cost keys. ON_DEMAND_FLEX (Vertex's flex opt-in
+    value) must map to "flex" so flex-tier requests are not billed as standard.
+    """
+    from litellm.cost_calculator import _map_traffic_type_to_service_tier
+
+    assert (
+        _map_traffic_type_to_service_tier(traffic_type) == expected_service_tier
+    )

tests/test_litellm/test_cost_calculator.py

@@ -2619,6 +2619,49 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(_local_mo
     assert cost == pytest.approx(expected_priority)
 
 
+def test_completion_cost_vertex_ai_gemini_flex_traffic_type(_local_model_cost_map):
+    """
+    Vertex AI flex-tier billing regression for issue #37647.
+
+    Vertex Gemini 3.x models route through ``cost_per_character`` (the
+    ``cost_router`` token-path gate only matches "gemini-2"), and its token
+    fallbacks dropped ``service_tier``. A response served with
+    ``trafficType=ON_DEMAND_FLEX`` must be billed at the flex rate, not the
+    standard rate.
+    """
+    from litellm import completion_cost
+
+    model = "gemini-3-test-flex-tier-cost-model"
+    litellm.register_model(
+        model_cost={
+            model: {
+                "input_cost_per_token": 1.5e-6,
+                "output_cost_per_token": 9e-6,
+                "input_cost_per_token_flex": 7.5e-7,
+                "output_cost_per_token_flex": 4.5e-6,
+                "litellm_provider": "vertex_ai",
+                "max_tokens": 8192,
+            }
+        }
+    )
+
+    def _cost_for_traffic_type(traffic_type):
+        usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
+        response = ModelResponse(usage=usage, model=model)
+        response._hidden_params["provider_specific_fields"] = {"traffic_type": traffic_type}
+        return completion_cost(
+            completion_response=response,
+            model=model,
+            custom_llm_provider="vertex_ai",
+        )
+
+    standard_cost = _cost_for_traffic_type("ON_DEMAND")
+    flex_cost = _cost_for_traffic_type("ON_DEMAND_FLEX")
+
+    assert standard_cost == pytest.approx(1000 * 1.5e-6 + 500 * 9e-6)
+    assert flex_cost == pytest.approx(1000 * 7.5e-7 + 500 * 4.5e-6)
+
+
 def test_completion_cost_non_string_service_tier_defers_to_served_tier(_local_model_cost_map):
     """
     Regression: a non-string request-level ``service_tier`` (reachable via