diff --git a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py index 7091563ae8f..836b02f2049 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py @@ -97,3 +97,43 @@ def calculate_tiered_cost( 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, +) -> Optional[dict]: + """ + Select the pricing tier for a request based on its total input token count. + + Alibaba Model Studio (Dashscope) tiered pricing is all-or-nothing: the tier is + chosen by the total input tokens of a single request and every token in the + request (input and output) is billed at that one tier's rate, rather than + graduated income-tax-style slicing. A tier matches when + ``range_start < input_tokens <= range_end`` (so a request of exactly + ``range_end`` tokens stays in the lower tier, matching the official + ``0 < Token <= 32K`` phrasing). Requests above the highest declared range fall + back to the last (most expensive) tier. + """ + if not tiered_pricing or input_tokens <= 0: + return None + + sorted_tiers = sorted(tiered_pricing, key=lambda t: t.get("range", [0, 0])[0]) + valid_tiers = [tier for tier in sorted_tiers if len(tier.get("range", [])) == 2] + if not valid_tiers: + return None + + matching = [tier for tier in valid_tiers if tier["range"][0] < input_tokens <= tier["range"][1]] + if matching: + return matching[0] + return valid_tiers[-1] + + +def tier_rate( + tier: dict, + cost_key: str, + fallback_cost_key: Optional[str] = None, +) -> float: + """Read a per-token rate from a tier, coercing YAML string costs to float.""" + raw = tier.get(cost_key) or tier.get(fallback_cost_key, 0) + return _coerce_cost_per_token(raw) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 55e9deaf125..f02b733229a 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -10,7 +10,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds -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.proxy._types import ( LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -925,9 +925,25 @@ def _estimate_request_input_cost_for_model( model: str, llm_router: Router | None, ) -> float | None: - model_info = _get_model_cost_info(model=model, llm_router=llm_router) - if model_info is None: - return None + estimates = [ + _input_cost_for_cost_info( + request_body=request_body, + route=route, + model=model, + model_info=model_info, + ) + for model_info in _get_model_cost_infos(model=model, llm_router=llm_router) + ] + valid_estimates = [estimate for estimate in estimates if estimate is not None] + return max(valid_estimates) if valid_estimates else None + + +def _input_cost_for_cost_info( + request_body: dict, + route: str, + model: str, + model_info: Dict[str, Any], +) -> Optional[float]: input_tokens = _estimate_input_tokens( request_body=request_body, route=route, @@ -938,11 +954,9 @@ def _estimate_request_input_cost_for_model( return None tiered_pricing = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: - return calculate_tiered_cost( - tokens=input_tokens, - tiered_pricing=tiered_pricing, - cost_key="input_cost_per_token", - ) + tier = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=input_tokens) + if tier is not None: + return input_tokens * tier_rate(tier, "input_cost_per_token") input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) if input_cost_per_token is None: return None @@ -955,10 +969,25 @@ def _estimate_request_max_cost_for_model( model: str, llm_router: Optional[Router], ) -> Optional[float]: - model_info = _get_model_cost_info(model=model, llm_router=llm_router) - if model_info is None: - return None + estimates = [ + _max_cost_for_cost_info( + request_body=request_body, + route=route, + model=model, + model_info=model_info, + ) + for model_info in _get_model_cost_infos(model=model, llm_router=llm_router) + ] + valid_estimates = [estimate for estimate in estimates if estimate is not None] + return max(valid_estimates) if valid_estimates else None + +def _max_cost_for_cost_info( + request_body: dict, + route: str, + model: str, + model_info: Dict[str, Any], +) -> Optional[float]: image_cost = _estimate_image_generation_cost( request_body=request_body, model_info=model_info, @@ -966,8 +995,6 @@ def _estimate_request_max_cost_for_model( if image_cost is not None: return image_cost - input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) - output_cost_per_token = _to_float(model_info.get("output_cost_per_token")) input_tokens = _estimate_input_tokens( request_body=request_body, route=route, @@ -985,16 +1012,14 @@ def _estimate_request_max_cost_for_model( output_multiplier = _get_output_multiplier(request_body=request_body) tiered_pricing = model_info.get("tiered_pricing") if isinstance(tiered_pricing, list) and tiered_pricing: - return calculate_tiered_cost( - tokens=input_tokens, - tiered_pricing=tiered_pricing, - cost_key="input_cost_per_token", - ) + calculate_tiered_cost( - tokens=output_tokens * output_multiplier, - tiered_pricing=tiered_pricing, - cost_key="output_cost_per_token", - ) + tier = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=input_tokens) + if tier is not None: + return (input_tokens * tier_rate(tier, "input_cost_per_token")) + ( + output_tokens * output_multiplier * tier_rate(tier, "output_cost_per_token") + ) + input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) + output_cost_per_token = _to_float(model_info.get("output_cost_per_token")) cost = 0.0 if input_cost_per_token is not None: cost += input_tokens * input_cost_per_token @@ -1052,46 +1077,69 @@ def _get_model_cost_info( llm_router: Optional[Router], ) -> Optional[Dict[str, Any]]: if llm_router is not None: - try: - model_group_info = llm_router.get_model_group_info(model_group=model) - if model_group_info is not None: - model_group_cost_info = model_group_info.model_dump() - # Reservation runs before routing, so the concrete deployment is - # unknown here. We assume deployments in a group share pricing and - # use the first tiered table we find as the estimate. This is a - # best-effort admission gate; post-request billing is exact against - # the deployment that actually served the request. Mixing tiered and - # flat deployments in one group can therefore over- or under-estimate - # at reservation time only. - deployments = llm_router.get_model_list(model_name=model) or [] - for deployment in deployments: - model_id = deployment.get("model_info", {}).get("id") - backend_model = deployment.get("litellm_params", {}).get("model") - if not isinstance(model_id, str) or not isinstance(backend_model, str): - continue - deployment_model_info = llm_router.get_deployment_model_info( - model_id=model_id, - model_name=backend_model, - ) - if deployment_model_info is None: - continue - tiered_pricing = deployment_model_info.get("tiered_pricing") - if isinstance(tiered_pricing, list) and tiered_pricing: - return { - **model_group_cost_info, - "tiered_pricing": tiered_pricing, - } - return model_group_cost_info - except Exception: - verbose_proxy_logger.debug( - "Unable to load router model group info for budget reservation", - exc_info=True, - ) + model_group_info = llm_router.get_model_group_info(model_group=model) + if model_group_info is not None: + return model_group_info.model_dump() + return dict(litellm.get_model_info(model=model)) + +def _get_model_cost_infos( + model: str, + llm_router: Optional[Router], +) -> List[Dict[str, Any]]: + """Cost-info candidates to estimate a request against for one model group. + + Reservation runs before routing, so the deployment that will serve the request + is unknown. Rather than guess, we estimate the cost against every eligible + pricing shape in the group (the group's flat rates plus each deployment's + tiered table) and let the caller reserve the maximum, so a cheaper sibling + deployment can never leave the request under-reserved. + """ try: - return dict(litellm.get_model_info(model=model)) + base = _get_model_cost_info(model=model, llm_router=llm_router) + if base is None: + return [] + tiered_tables = _get_deployment_tiered_pricing_tables(model=model, llm_router=llm_router) except Exception: + verbose_proxy_logger.debug( + "Unable to load model cost info for budget reservation", + exc_info=True, + ) + return [] + if not tiered_tables: + return [base] + return [base, *({**base, "tiered_pricing": table} for table in tiered_tables)] + + +def _deployment_tiered_pricing_table( + deployment: Dict[str, Any], + llm_router: Router, +) -> Optional[List[dict]]: + model_id = deployment.get("model_info", {}).get("id") + backend_model = deployment.get("litellm_params", {}).get("model") + if not isinstance(model_id, str) or not isinstance(backend_model, str): return None + deployment_model_info = llm_router.get_deployment_model_info(model_id=model_id, model_name=backend_model) + if deployment_model_info is None: + return None + tiered_pricing = deployment_model_info.get("tiered_pricing") + if isinstance(tiered_pricing, list) and tiered_pricing: + return tiered_pricing + return None + + +def _get_deployment_tiered_pricing_tables( + model: str, + llm_router: Optional[Router], +) -> List[List[dict]]: + if llm_router is None: + return [] + deployments = llm_router.get_model_list(model_name=model) or [] + return [ + table + for deployment in deployments + if (table := _deployment_tiered_pricing_table(deployment, llm_router)) is not None + ] def _estimate_input_tokens( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 8e00db78a08..efc4c9758bf 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -824,6 +824,90 @@ async def test_should_reserve_tiered_pricing_cost(spend_counter_state): await release_budget_reservation(reservation) +def test_tiered_reservation_is_all_or_nothing_with_output_tier_from_input_length(): + """Dashscope tiered pricing is all-or-nothing: the tier is chosen by the total + input tokens and every token (input and output) is billed at that tier's rate. + + A long-context request with a large output allowance must reserve the output at + the input-selected tier, not at the cheapest tier picked from the output volume. + The earlier graduated calculation under-reserved such requests, letting a caller + slip past a depleted budget.""" + tiered_pricing = [ + {"range": [0, 32000], "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}, + {"range": [32000, 128000], "input_cost_per_token": 4e-06, "output_cost_per_token": 8e-06}, + ] + input_tokens = 100000 # falls entirely in the second tier + output_tokens = 1000 + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={"tiered_pricing": tiered_pricing, "max_output_tokens": 200000}, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_input_tokens", + return_value=input_tokens, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_output_tokens", + return_value=output_tokens, + ), + ): + estimated = estimate_request_max_cost( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + ) + + expected = (input_tokens * 4e-06) + (output_tokens * 8e-06) + assert estimated == pytest.approx(expected) + + # What the old graduated math (with the output tier taken from output volume) + # would have reserved. The all-or-nothing estimate must be strictly larger. + graduated_under_reserve = (32000 * 1e-06) + (68000 * 4e-06) + (output_tokens * 2e-06) + assert estimated > graduated_under_reserve + + +def test_reservation_uses_most_expensive_deployment_in_group(): + """When a model group mixes deployments with different tiered rates, reservation + must estimate against the most expensive one. Reserving the cheaper sibling would + let a caller repeatedly hit the alias and exceed the budget once routed to the + costlier deployment.""" + cheap = [{"range": [0, 32000], "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}] + expensive = [{"range": [0, 32000], "input_cost_per_token": 5e-06, "output_cost_per_token": 1e-05}] + input_tokens = 1000 + output_tokens = 10 + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={"max_output_tokens": 200000}, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._get_deployment_tiered_pricing_tables", + return_value=[cheap, expensive], + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_input_tokens", + return_value=input_tokens, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_output_tokens", + return_value=output_tokens, + ), + ): + estimated = estimate_request_max_cost( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + ) + + expected_expensive = (input_tokens * 5e-06) + (output_tokens * 1e-05) + expected_cheap = (input_tokens * 1e-06) + (output_tokens * 2e-06) + assert expected_expensive > expected_cheap + assert estimated == pytest.approx(expected_expensive) + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequests( spend_counter_state,