fix(dashscope): bill tiered pricing at the request-size tier

This commit is contained in:
Devin AI 2026-07-26 22:07:59 +00:00
parent 24123269cc
commit 66317ce63d
3 changed files with 105 additions and 178 deletions

View file

@ -1,5 +1,5 @@
"""
Provider-neutral graduated tiered pricing calculation.
Provider-neutral request-size tiered pricing helpers.
Shared by provider cost calculators (e.g. Dashscope) and the proxy budget
reservation logic so neither has to depend on the other.
@ -25,80 +25,6 @@ def _coerce_cost_per_token(value: Union[float, int, str, None]) -> float:
return float(value)
def calculate_tiered_cost(
tokens: int,
tiered_pricing: List[dict],
cost_key: str,
fallback_cost_key: Optional[str] = None,
) -> float:
"""
Calculate cost for a given number of tokens based on a true tiered pricing structure.
This function iterates through sorted pricing tiers, calculates the cost for the
number of tokens that fall into each tier's range, and sums them up to get the total cost.
Args:
tokens (int): The total number of tokens to calculate the cost for.
tiered_pricing (List[dict]): A list of dictionaries, where each dictionary
represents a pricing tier.
cost_key (str): The key in the tier dictionary that holds the per-token cost
(e.g., 'input_cost_per_token').
fallback_cost_key (Optional[str], optional): A fallback key to use if the
primary `cost_key` is not found in a tier. Defaults to None.
Returns:
float: The total calculated cost for the given tokens.
Example:
>>> tiered_pricing = [
... {"range": [0, 100000], "input_cost_per_token": 0.0001},
... {"range": [100000, 500000], "input_cost_per_token": 0.00005},
... ]
Calculating cost for 150,000 tokens:
(100,000 * 0.0001) + (50,000 * 0.00005) = $12.5
"""
if not tiered_pricing or tokens <= 0:
return 0.0
total_cost = 0.0
tokens_processed = 0
sorted_tiers = sorted(tiered_pricing, key=lambda x: x.get("range", [0, 0])[0])
for tier in sorted_tiers:
if tokens_processed >= tokens:
break
tier_range = tier.get("range", [])
if len(tier_range) != 2:
continue
range_start, range_end = tier_range
if tokens <= range_start:
continue
tier_start = max(range_start, tokens_processed)
tier_end = min(range_end, tokens)
if tier_end > tier_start:
tokens_in_tier = tier_end - tier_start
cost_per_token = tier.get(cost_key) or tier.get(fallback_cost_key, 0)
total_cost += tokens_in_tier * _coerce_cost_per_token(cost_per_token)
tokens_processed = tier_end
# After loop, check if any tokens remain (i.e., tokens > highest tier's end range)
# and charge them at the last tier's rate.
if tokens_processed < tokens and sorted_tiers:
last_tier = sorted_tiers[-1]
remaining_tokens = tokens - tokens_processed
cost_per_token = last_tier.get(cost_key) or last_tier.get(fallback_cost_key, 0)
total_cost += remaining_tokens * _coerce_cost_per_token(cost_per_token)
return total_cost
def select_tier_for_input(
tiered_pricing: List[dict],
input_tokens: int,

View file

@ -7,7 +7,7 @@ Handles tiered pricing and prompt caching scenarios.
from dataclasses import dataclass
from typing import List, Optional, Tuple
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import calculate_tiered_cost
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import get_model_info
@ -46,22 +46,13 @@ def _extract_token_breakdown(usage: Usage) -> TokenBreakdown:
def _calculate_prompt_cost(
breakdown: TokenBreakdown,
model_info: ModelInfo,
tiered_pricing: Optional[List[dict]],
tier: Optional[dict],
) -> float:
"""Calculate total prompt cost including cached tokens."""
if tiered_pricing:
text_cost = calculate_tiered_cost(
tokens=breakdown.text_tokens,
tiered_pricing=tiered_pricing,
cost_key="input_cost_per_token",
)
cache_cost = calculate_tiered_cost(
tokens=breakdown.cached_tokens,
tiered_pricing=tiered_pricing,
cost_key="cache_read_input_token_cost",
fallback_cost_key="input_cost_per_token",
)
return text_cost + cache_cost
if tier is not None:
input_cost = tier_rate(tier, "input_cost_per_token")
cache_cost = tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token")
return (breakdown.text_tokens * input_cost) + (breakdown.cached_tokens * cache_cost)
input_cost = float(model_info.get("input_cost_per_token") or 0.0)
@ -78,22 +69,13 @@ def _calculate_prompt_cost(
def _calculate_completion_cost(
breakdown: TokenBreakdown,
model_info: ModelInfo,
tiered_pricing: Optional[List[dict]],
tier: Optional[dict],
) -> float:
"""Calculate total completion cost including reasoning tokens."""
if tiered_pricing:
completion_cost = calculate_tiered_cost(
tokens=breakdown.completion_tokens,
tiered_pricing=tiered_pricing,
cost_key="output_cost_per_token",
)
reasoning_cost = calculate_tiered_cost(
tokens=breakdown.reasoning_tokens,
tiered_pricing=tiered_pricing,
cost_key="output_cost_per_reasoning_token",
fallback_cost_key="output_cost_per_token",
)
return completion_cost + reasoning_cost
if tier is not None:
output_cost = tier_rate(tier, "output_cost_per_token")
reasoning_cost = tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token")
return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost)
output_cost = float(model_info.get("output_cost_per_token") or 0.0)
@ -111,7 +93,10 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]:
"""
Calculate cost per token for Dashscope models.
Supports both tiered and flat pricing with cached and reasoning tokens.
Alibaba Model Studio tiered pricing selects a single tier from the request's
total input token count and bills every token of the request at that tier's
rates, so the tier is resolved once here and shared by the prompt and
completion legs.
Args:
model: Model name without provider prefix
@ -122,11 +107,16 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]:
"""
model_info = get_model_info(model=model, custom_llm_provider="dashscope")
breakdown = _extract_token_breakdown(usage)
tiered_pricing = model_info.get("tiered_pricing") if isinstance(model_info.get("tiered_pricing"), list) else None
raw_tiered_pricing = model_info.get("tiered_pricing")
tiered_pricing: Optional[List[dict]] = raw_tiered_pricing if isinstance(raw_tiered_pricing, list) else None
prompt_cost = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing)
completion_cost = _calculate_completion_cost(
breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing
tier = (
select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=usage.prompt_tokens or 0)
if tiered_pricing
else None
)
prompt_cost = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier)
completion_cost = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier)
return prompt_cost, completion_cost

View file

@ -2,13 +2,12 @@
Test suite for Dashscope cost calculation functionality.
Tests the cost calculation for Dashscope models including:
- Correctly calculates graduated tiered pricing.
- Selects a single pricing tier from the request's total input size.
- Falls back to flat-rate pricing for non-tiered models.
- Handles interactions with cached tokens.
- Correctly calculates costs for token counts exceeding the highest defined tier.
"""
import json
import math
import os
import sys
@ -55,7 +54,7 @@ class TestDashscopeCostCalculator:
def test_dashscope_tiered_pricing_within_first_tier(self):
"""
Tests the dashscope tiered pricing when token count is entirely within the first tier.
Tests the dashscope tiered pricing when the request's input size falls in the first tier.
Uses 'dashscope/qwen-flash' as a real-world example.
"""
# Tier 1 for qwen-flash is [0, 256,000] tokens
@ -73,13 +72,15 @@ class TestDashscopeCostCalculator:
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
def test_dashscope_tiered_pricing_spanning_multiple_tiers(self):
def test_dashscope_tier_selected_by_total_input_size(self):
"""
Tests the dashscope tiered pricing with the corrected graduated calculation logic.
This is the most important test for validating the fix.
Regression for graduated slicing: Alibaba Model Studio picks one tier from the
request's total input tokens and bills every input and output token at that
tier's rates. A 300k-input request must be billed entirely at tier 2, not
sliced across tier 1 and tier 2.
"""
# Tiering for qwen-flash: Tier 1: [0, 256k], Tier 2: [256k, 1M]
usage = Usage(prompt_tokens=300000, completion_tokens=300000)
usage = Usage(prompt_tokens=300000, completion_tokens=2000)
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen-flash", usage=usage
)
@ -88,23 +89,41 @@ class TestDashscopeCostCalculator:
tier_1 = model_info["tiered_pricing"][0]
tier_2 = model_info["tiered_pricing"][1]
# Expected prompt cost: (256,000 tokens * tier_1_price) + (44,000 tokens * tier_2_price)
expected_prompt_cost = (256000 * tier_1["input_cost_per_token"]) + (
44000 * tier_2["input_cost_per_token"]
)
# Expected completion cost: (256,000 tokens * tier_1_price) + (44,000 tokens * tier_2_price)
expected_completion_cost = (256000 * tier_1["output_cost_per_token"]) + (
44000 * tier_2["output_cost_per_token"]
)
expected_prompt_cost = 300000 * tier_2["input_cost_per_token"]
expected_completion_cost = 2000 * tier_2["output_cost_per_token"]
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
graduated_prompt_cost = (256000 * tier_1["input_cost_per_token"]) + (
44000 * tier_2["input_cost_per_token"]
)
assert prompt_cost > graduated_prompt_cost
def test_dashscope_tier_boundary_stays_in_lower_tier(self):
"""
A request whose input size exactly equals a tier's upper bound belongs to that
tier (the docs phrase ranges as ``0 < Token <= 256K``).
"""
usage = Usage(prompt_tokens=256000, completion_tokens=1000)
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen-flash", usage=usage
)
tier_1 = litellm.get_model_info("dashscope/qwen-flash")["tiered_pricing"][0]
assert math.isclose(
prompt_cost, 256000 * tier_1["input_cost_per_token"], rel_tol=1e-10
)
assert math.isclose(
completion_cost, 1000 * tier_1["output_cost_per_token"], rel_tol=1e-10
)
def test_dashscope_tiered_pricing_with_caching(self):
"""
Tests tiered pricing with cached tokens. This replaces the old, incorrect test.
Uses qwen3-coder-plus, which has cache-specific pricing defined.
Cached tokens count toward the request's input size for tier selection and are
billed at the selected tier's cache rate, rather than being tiered from zero
on their own.
"""
usage = Usage(
prompt_tokens=50000, # 10k cached + 40k new
@ -116,20 +135,14 @@ class TestDashscopeCostCalculator:
prompt_cost, _ = dashscope_cost_per_token(model="qwen3-coder-plus", usage=usage)
model_info = litellm.get_model_info("dashscope/qwen3-coder-plus")
tier_1 = model_info["tiered_pricing"][0]
# 50k total input selects tier 2 ([32k, 128k])
tier_2 = model_info["tiered_pricing"][1]
# 10k cached tokens are all in the first tier
expected_cache_cost = 10000 * tier_1["cache_read_input_token_cost"]
# 40k new tokens: 32k in tier 1, and the remaining 8k in tier 2
expected_text_cost = (32000 * tier_1["input_cost_per_token"]) + (
8000 * tier_2["input_cost_per_token"]
expected_prompt_cost = (10000 * tier_2["cache_read_input_token_cost"]) + (
40000 * tier_2["input_cost_per_token"]
)
expected_total_prompt_cost = expected_cache_cost + expected_text_cost
assert math.isclose(prompt_cost, expected_total_prompt_cost, rel_tol=1e-10)
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
def _register_string_valued_tiered_model(self, model_key: str) -> None:
"""Register a model whose tier costs are strings, mimicking YAML config parsing."""
@ -152,8 +165,8 @@ class TestDashscopeCostCalculator:
def test_dashscope_tiered_pricing_string_costs_within_tier(self):
"""
Regression: YAML-parsed tier costs can be strings (e.g. "4e-07"). Costs that
fall entirely within a single tier must still be computed as floats.
Regression: YAML-parsed tier costs can be strings (e.g. "4e-07") and must still
be computed as floats.
"""
self._register_string_valued_tiered_model("dashscope/qwen-str-tier-test")
@ -162,18 +175,13 @@ class TestDashscopeCostCalculator:
model="qwen-str-tier-test", usage=usage
)
expected_prompt_cost = 500 * float("4e-07")
expected_completion_cost = 200 * float("1.6e-06")
assert prompt_cost > 0
assert completion_cost > 0
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
assert math.isclose(prompt_cost, 500 * float("4e-07"), rel_tol=1e-10)
assert math.isclose(completion_cost, 200 * float("1.6e-06"), rel_tol=1e-10)
def test_dashscope_tiered_pricing_string_costs_exceeding_highest_tier(self):
"""
Regression: string-valued tier costs must also be coerced in the
remaining-tokens path that charges tokens above the highest tier.
Requests larger than the highest declared range fall back to the last tier, and
string-valued costs there must also be coerced.
"""
self._register_string_valued_tiered_model("dashscope/qwen-str-tier-test")
@ -182,43 +190,46 @@ class TestDashscopeCostCalculator:
model="qwen-str-tier-test", usage=usage
)
# prompt: 1000 @ tier1 + 1000 @ tier2 + 500 remaining @ tier2 rate
expected_prompt_cost = (
(1000 * float("4e-07")) + (1000 * float("8e-07")) + (500 * float("8e-07"))
)
# completion: 1000 @ tier1 + 1000 @ tier2 + 1000 remaining @ tier2 rate
expected_completion_cost = (
(1000 * float("1.6e-06")) + (1000 * float("3.2e-06")) + (1000 * float("3.2e-06"))
)
assert prompt_cost > 0
assert completion_cost > 0
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)
assert math.isclose(prompt_cost, 2500 * float("8e-07"), rel_tol=1e-10)
assert math.isclose(completion_cost, 3000 * float("3.2e-06"), rel_tol=1e-10)
def test_dashscope_tiered_pricing_exceeding_highest_tier(self):
"""
Tests tiered pricing when token count exceeds the highest defined tier range.
This replaces the old, incorrect test and validates the new fallback logic.
Tests tiered pricing when the input size exceeds the highest defined tier range;
the most expensive tier applies to the whole request.
"""
usage = Usage(
prompt_tokens=1200000, completion_tokens=1000
) # Max defined range for qwen-flash is 1M
prompt_cost, _ = dashscope_cost_per_token(model="qwen-flash", usage=usage)
model_info = litellm.get_model_info("dashscope/qwen-flash")
tier_1 = model_info["tiered_pricing"][0]
tier_2 = model_info["tiered_pricing"][1]
# Expected cost: (tier_1_tokens * tier_1_price) + (tokens_up_to_max_range_in_tier_2 * tier_2_price) + (remaining_tokens * tier_2_price)
tokens_in_tier_2_range = 1000000 - 256000
remaining_tokens_over_max = 1200000 - 1000000
expected_prompt_cost = (
(256000 * tier_1["input_cost_per_token"])
+ (tokens_in_tier_2_range * tier_2["input_cost_per_token"])
+ (remaining_tokens_over_max * tier_2["input_cost_per_token"])
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen-flash", usage=usage
)
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
tier_2 = litellm.get_model_info("dashscope/qwen-flash")["tiered_pricing"][1]
assert math.isclose(
prompt_cost, 1200000 * tier_2["input_cost_per_token"], rel_tol=1e-10
)
assert math.isclose(
completion_cost, 1000 * tier_2["output_cost_per_token"], rel_tol=1e-10
)
def test_dashscope_cost_matches_budget_reservation_estimate(self):
"""
The post-response spend must agree with the proxy's pre-call reservation
estimate, which selects a tier by request input size.
"""
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import (
select_tier_for_input,
tier_rate,
)
usage = Usage(prompt_tokens=300000, completion_tokens=2000)
prompt_cost, _ = dashscope_cost_per_token(model="qwen-flash", usage=usage)
tiered_pricing = litellm.get_model_info("dashscope/qwen-flash")["tiered_pricing"]
tier = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=300000)
reserved_input_cost = 300000 * tier_rate(tier, "input_cost_per_token")
assert math.isclose(prompt_cost, reserved_input_cost, rel_tol=1e-10)