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.
This commit is contained in:
mateo-berri 2026-08-26 16:29:45 -07:00
parent 909ab23b89
commit ac3f987883
3 changed files with 54 additions and 0 deletions

View file

@ -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":

View file

@ -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)

View file

@ -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