mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
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:
parent
909ab23b89
commit
ac3f987883
3 changed files with 54 additions and 0 deletions
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue