feat: support dashscope tiered pricing

* add helper functions

* update generic_cost_per_token function

* add test

* formatting

* add examples in docstring for _calculate_tiered_cost

* Restore files to upstream/main version

* dashscope specific calculation

* improve for different costs

* remove _calculate_flat_cost function
This commit is contained in:
Alan Ponnachan 2025-11-04 09:57:25 +05:30 committed by GitHub
parent 5ddc5410bf
commit 45cf6fc289
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 242 additions and 204 deletions

View file

@ -14,6 +14,7 @@ from litellm.utils import get_model_info
@dataclass
class TokenBreakdown:
"""Token breakdown for cost calculation."""
text_tokens: int
cached_tokens: int
completion_tokens: int
@ -23,133 +24,194 @@ class TokenBreakdown:
def _extract_token_breakdown(usage: Usage) -> TokenBreakdown:
"""Extract token counts from usage, handling cached and reasoning tokens."""
cached_tokens = 0
if usage.prompt_tokens_details and hasattr(usage.prompt_tokens_details, "cached_tokens"):
if usage.prompt_tokens_details and hasattr(
usage.prompt_tokens_details, "cached_tokens"
):
cached_tokens = usage.prompt_tokens_details.cached_tokens or 0
text_tokens = usage.prompt_tokens - cached_tokens
reasoning_tokens = 0
if (hasattr(usage, "completion_tokens_details") and
usage.completion_tokens_details and
hasattr(usage.completion_tokens_details, "reasoning_tokens")):
if (
hasattr(usage, "completion_tokens_details")
and usage.completion_tokens_details
and hasattr(usage.completion_tokens_details, "reasoning_tokens")
):
reasoning_tokens = usage.completion_tokens_details.reasoning_tokens or 0
completion_tokens = (usage.completion_tokens or 0) - reasoning_tokens
return TokenBreakdown(text_tokens, cached_tokens, completion_tokens, reasoning_tokens)
return TokenBreakdown(
text_tokens, cached_tokens, completion_tokens, reasoning_tokens
)
def _calculate_tiered_cost(
tokens: int,
tiered_pricing: List[dict],
tokens: int,
tiered_pricing: List[dict],
cost_key: str,
fallback_cost_key: Optional[str] = None
fallback_cost_key: Optional[str] = None,
) -> float:
"""Calculate cost using tiered pricing structure.
Finds the appropriate tier based on token count and applies that tier's rate to all tokens.
"""
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
# Find the appropriate tier for the token count
for tier in tiered_pricing:
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
# Check if tokens fall within this tier's range
if range_start <= tokens <= range_end:
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)
return tokens * cost_per_token
# If no tier matches, use the last tier (highest tier)
if tiered_pricing:
last_tier = tiered_pricing[-1]
total_cost += tokens_in_tier * 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)
return tokens * cost_per_token
return 0.0
total_cost += remaining_tokens * cost_per_token
return total_cost
def _calculate_flat_cost(tokens: int, cost_per_token: float) -> float:
"""Calculate cost using flat pricing."""
return tokens * cost_per_token
def _calculate_prompt_cost(breakdown: TokenBreakdown, model_info: ModelInfo, tiered_pricing: Optional[List[dict]]) -> float:
def _calculate_prompt_cost(
breakdown: TokenBreakdown,
model_info: ModelInfo,
tiered_pricing: Optional[List[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"
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"
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
input_cost = model_info.get("input_cost_per_token", 0.0)
cache_cost = model_info.get("cache_read_input_token_cost", input_cost) or input_cost
return (_calculate_flat_cost(tokens=breakdown.text_tokens, cost_per_token=input_cost) +
_calculate_flat_cost(tokens=breakdown.cached_tokens, cost_per_token=cache_cost))
input_cost = float(model_info.get("input_cost_per_token") or 0.0)
# For cache_cost, first try the specific key, then fall back to input_cost.
cache_cost_val = model_info.get("cache_read_input_token_cost")
if cache_cost_val is None:
cache_cost = input_cost
else:
cache_cost = float(cache_cost_val)
return (breakdown.text_tokens * input_cost) + (breakdown.cached_tokens * cache_cost)
def _calculate_completion_cost(breakdown: TokenBreakdown, model_info: ModelInfo, tiered_pricing: Optional[List[dict]]) -> float:
def _calculate_completion_cost(
breakdown: TokenBreakdown,
model_info: ModelInfo,
tiered_pricing: Optional[List[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"
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,
tokens=breakdown.reasoning_tokens,
tiered_pricing=tiered_pricing,
cost_key="output_cost_per_reasoning_token",
fallback_cost_key="output_cost_per_token"
fallback_cost_key="output_cost_per_token",
)
return completion_cost + reasoning_cost
output_cost = model_info.get("output_cost_per_token", 0.0)
reasoning_cost = model_info.get("output_cost_per_reasoning_token", output_cost) or output_cost
return (_calculate_flat_cost(tokens=breakdown.completion_tokens, cost_per_token=output_cost) +
_calculate_flat_cost(tokens=breakdown.reasoning_tokens, cost_per_token=reasoning_cost))
output_cost = float(model_info.get("output_cost_per_token") or 0.0)
# For reasoning_cost, first try the specific key, then fall back to output_cost.
reasoning_cost_val = model_info.get("output_cost_per_reasoning_token")
if reasoning_cost_val is None:
reasoning_cost = output_cost
else:
reasoning_cost = float(reasoning_cost_val)
return (breakdown.completion_tokens * output_cost) + (
breakdown.reasoning_tokens * reasoning_cost
)
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.
Args:
model: Model name without provider prefix
usage: LiteLLM Usage block
Returns:
Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd)
"""
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
tiered_pricing = (
model_info.get("tiered_pricing")
if isinstance(model_info.get("tiered_pricing"), list)
else None
)
prompt_cost = _calculate_prompt_cost(
breakdown=breakdown,
model_info=model_info,
tiered_pricing=tiered_pricing
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
breakdown=breakdown, model_info=model_info, tiered_pricing=tiered_pricing
)
return prompt_cost, completion_cost

View file

@ -2,10 +2,10 @@
Test suite for Dashscope cost calculation functionality.
Tests the cost calculation for Dashscope models including:
- Tiered pricing based on input token ranges
- Caching discounts
- Reasoning tokens
- Standard flat pricing fallback
- Correctly calculates graduated tiered pricing.
- 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
@ -22,11 +22,7 @@ import litellm
from litellm.llms.dashscope.cost_calculator import (
cost_per_token as dashscope_cost_per_token,
)
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
PromptTokensDetailsWrapper,
Usage,
)
from litellm.types.utils import Usage, PromptTokensDetailsWrapper
class TestDashscopeCostCalculator:
@ -34,150 +30,130 @@ class TestDashscopeCostCalculator:
@pytest.fixture(autouse=True)
def setup_model_cost_map(self):
"""Set up the model cost map for testing."""
# Ensure we use local model cost map for consistent testing
"""Set up the model cost map for testing by loading it locally."""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
# Find the project root directory and load model cost map
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = current_dir
while not os.path.exists(os.path.join(project_root, "model_prices_and_context_window.json")):
parent = os.path.dirname(project_root)
if parent == project_root: # Reached filesystem root
break
project_root = parent
model_cost_path = os.path.join(project_root, "model_prices_and_context_window.json")
with open(model_cost_path, "r") as f:
model_cost_map = json.load(f)
litellm.model_cost = model_cost_map
litellm.model_cost = litellm.get_model_cost_map(url="")
def test_flat_pricing_basic_cost_calculation(self):
"""Test basic cost calculation for flat pricing models (qwen-max)."""
usage = Usage(
prompt_tokens=1000,
completion_tokens=500,
total_tokens=1500
)
def test_dashscope_flat_pricing_fallback(self):
"""
Tests that the dashscope calculator falls back to flat pricing for models
without a 'tiered_pricing' key (e.g., qwen-max).
"""
usage = Usage(prompt_tokens=1000, completion_tokens=500)
# We call the specific calculator for dashscope
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen-max",
usage=usage
model="qwen-max", usage=usage
)
# Expected costs for qwen-max:
# Input: 1000 tokens * $1.6e-6 = $0.0016
# Output: 500 tokens * $6.4e-6 = $0.0032
expected_prompt_cost = 1000 * 1.6e-6
expected_completion_cost = 500 * 6.4e-6
model_info = litellm.get_model_info("dashscope/qwen-max")
expected_prompt_cost = 1000 * model_info["input_cost_per_token"]
expected_completion_cost = 500 * model_info["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)
def test_tiered_pricing_single_tier(self):
"""Test tiered pricing when all tokens fall within first tier."""
usage = Usage(
prompt_tokens=20000, # Within first tier (0-32K)
completion_tokens=1000,
total_tokens=21000
)
def test_dashscope_tiered_pricing_within_first_tier(self):
"""
Tests the dashscope tiered pricing when token count is entirely within the first tier.
Uses 'dashscope/qwen-flash' as a real-world example.
"""
# Tier 1 for qwen-flash is [0, 256,000] tokens
usage = Usage(prompt_tokens=100000, completion_tokens=50000)
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen3-coder-plus",
usage=usage
model="qwen-flash", usage=usage
)
# Expected costs for qwen3-coder-plus (tier 1):
# Input: 20,000 tokens * $1e-6 = $0.02
# Output: 1,000 tokens * $5e-6 = $0.005
expected_prompt_cost = 20000 * 1e-6
expected_completion_cost = 1000 * 5e-6
model_info = litellm.get_model_info("dashscope/qwen-flash")
tier_1_pricing = model_info["tiered_pricing"][0]
expected_prompt_cost = 100000 * tier_1_pricing["input_cost_per_token"]
expected_completion_cost = 50000 * tier_1_pricing["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)
def test_tiered_pricing_higher_tier(self):
"""Test tiered pricing when tokens fall in higher tier (tier 3)."""
usage = Usage(
prompt_tokens=150000, # Falls in tier 3 (128K-256K)
completion_tokens=2000,
total_tokens=152000
)
def test_dashscope_tiered_pricing_spanning_multiple_tiers(self):
"""
Tests the dashscope tiered pricing with the corrected graduated calculation logic.
This is the most important test for validating the fix.
"""
# Tiering for qwen-flash: Tier 1: [0, 256k], Tier 2: [256k, 1M]
usage = Usage(prompt_tokens=300000, completion_tokens=300000)
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen3-coder-plus",
usage=usage
model="qwen-flash", usage=usage
)
# Expected input cost calculation:
# 150,000 tokens falls in tier 3 (128K-256K), so all tokens are charged at tier 3 rate
# Input: 150,000 tokens * $3e-6 = $0.45
# Output: 2,000 tokens falls in tier 1 (0-32K), so charged at tier 1 rate
# Output: 2,000 tokens * $5e-6 = $0.01
expected_prompt_cost = 150000 * 3e-6 # All tokens at tier 3 rate
expected_completion_cost = 2000 * 5e-6 # All tokens at tier 1 rate
model_info = litellm.get_model_info("dashscope/qwen-flash")
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"]
)
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_tiered_pricing_with_caching(self):
"""Test tiered pricing with cached tokens."""
prompt_tokens_details = PromptTokensDetailsWrapper(
cached_tokens=10000 # 10K cached tokens
)
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.
"""
usage = Usage(
prompt_tokens=50000, # 40K regular + 10K cached = 50K total
prompt_tokens=50000, # 10k cached + 40k new
completion_tokens=1000,
total_tokens=51000,
prompt_tokens_details=prompt_tokens_details
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=10000),
)
prompt_cost, completion_cost = dashscope_cost_per_token(
model="qwen3-coder-plus",
usage=usage
)
# Expected cost calculation:
# Regular tokens: 40,000 falls in tier 2 (32K-128K), so all charged at tier 2 rate
# - Regular: 40,000 * $1.8e-6 = $0.072
# Cached tokens: 10,000 falls in tier 1 (0-32K), so charged at tier 1 cached rate
# - Cached: 10,000 * $1e-7 = $0.001
# Total input cost = $0.072 + $0.001 = $0.073
regular_tokens = 40000
cached_tokens = 10000
expected_regular_cost = regular_tokens * 1.8e-6 # Tier 2 rate
expected_cached_cost = cached_tokens * 1e-7 # Tier 1 cached rate
expected_prompt_cost = expected_regular_cost + expected_cached_cost
expected_completion_cost = 1000 * 5e-6 # Tier 1 rate
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_tiered_pricing_highest_tier(self):
"""Test tiered pricing when tokens exceed highest tier range."""
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]
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_total_prompt_cost = expected_cache_cost + expected_text_cost
assert math.isclose(prompt_cost, expected_total_prompt_cost, 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.
"""
usage = Usage(
prompt_tokens=2000000, # Exceeds tier 4 max (1M), should use tier 4 rate
completion_tokens=5000,
total_tokens=2005000
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="qwen3-coder-plus",
usage=usage
)
# Expected cost calculation:
# 2,000,000 tokens exceeds tier 4 (256K-1M), so use tier 4 rate for all tokens
# Input: 2,000,000 tokens * $6e-6 = $12.0
# Output: 5,000 tokens falls in tier 1 (0-32K), so charged at tier 1 rate
# Output: 5,000 tokens * $5e-6 = $0.025
expected_prompt_cost = 2000000 * 6e-6 # Tier 4 rate (highest tier)
expected_completion_cost = 5000 * 5e-6 # Tier 1 rate
assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10)
assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10)