From ac3f9878837f89e8b21c1bf86a642040ff74cff2 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 26 Aug 2026 16:29:45 -0700
Subject: [PATCH] fix: thread service_tier through vertex cost_per_character
 fallbacks

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, so ON_DEMAND_FLEX responses were still billed at the standard
rate. Pass the tier through the call site and all four fallbacks.
---
 litellm/cost_calculator.py                 |  1 +
 litellm/llms/vertex_ai/cost_calculator.py  |  7 ++++
 tests/test_litellm/test_cost_calculator.py | 46 ++++++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 46a42b616f8..b8ffdc1bf74 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -589,6 +589,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":
diff --git a/litellm/llms/vertex_ai/cost_calculator.py b/litellm/llms/vertex_ai/cost_calculator.py
index 23cb1e5b580..8b00fc2e925 100644
--- a/litellm/llms/vertex_ai/cost_calculator.py
+++ b/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)
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 98938dee62e..f0a84fe620b 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -2613,6 +2613,52 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate():
     assert cost == pytest.approx(expected_priority)
 
 
+def test_completion_cost_vertex_ai_gemini_flex_traffic_type(monkeypatch):
+    """
+    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
+
+    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
+
+    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():
     """
     Regression: a non-string request-level ``service_tier`` (reachable via