From 6d432966f43a81f03471384c65182f39cbb12879 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97SO?= <142557582+Linxiushen@users.noreply.github.com> Date: Mon, 28 Sep 2026 01:14:31 +0800 Subject: [PATCH] fix(router): resolve provider-prefixed model costs --- litellm/router_strategy/lowest_cost.py | 45 ++++++++++++++++++- .../unit/router_strategy/test_lowest_cost.py | 45 +++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 5894da9cb4f..6c05029cb76 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -4,6 +4,9 @@ from collections.abc import Mapping from datetime import datetime from typing import Final +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm import ModelResponse, token_counter, verbose_logger from litellm._internal_context import with_service_target @@ -11,14 +14,14 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.batch_utils import is_batch_retrieve_call_type -from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed +from litellm.router_utils.deployment_pricing import deployment_pricing_models def _deployment_cost_map_rates(deployment: object) -> tuple[float | None, float | None]: """Input and output per-token prices from the first cost-map row among the deployment's pricing models that carries one; (None, None) when no row does, so the caller falls back to its unmapped default.""" - rows: Final = (str_keyed(_cost_map().get(model)) for model in deployment_pricing_models(deployment)) + rows: Final = (_get_model_cost_info(model) for model in deployment_pricing_models(deployment)) row: Final = next((row for row in rows if _declares_token_price(row)), None) if row is None: return None, None @@ -37,6 +40,44 @@ def _as_rate(value: object) -> float | None: return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None +class _ModelCostInfo(TypedDict, total=False): + input_cost_per_token: ReadOnly[float | None] + output_cost_per_token: ReadOnly[float | None] + litellm_provider: ReadOnly[str | None] + + +_MODEL_COST_INFO_ADAPTER: Final = TypeAdapter(_ModelCostInfo) + + +def _get_validated_model_cost_info(model_name: str) -> _ModelCostInfo | None: + raw_model_cost_info: Final[object] = _cost_map().get(model_name) + if raw_model_cost_info is None: + return None + + try: + return _MODEL_COST_INFO_ADAPTER.validate_python(raw_model_cost_info) + except ValidationError: + return None + + +def _get_model_cost_info(model_name: str | None) -> _ModelCostInfo: + if model_name is None: + return {} + + exact_model_cost: Final = _get_validated_model_cost_info(model_name) + if exact_model_cost is not None: + return exact_model_cost + + provider_name, separator, unprefixed_model_name = model_name.partition("/") + if separator == "": + return {} + + unprefixed_model_cost: Final = _get_validated_model_cost_info(unprefixed_model_name) + if unprefixed_model_cost is None or unprefixed_model_cost.get("litellm_provider") != provider_name: + return {} + return unprefixed_model_cost + + class LowestCostLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 diff --git a/tests/unit/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py index e3fdcd69499..c3190d0e359 100644 --- a/tests/unit/router_strategy/test_lowest_cost.py +++ b/tests/unit/router_strategy/test_lowest_cost.py @@ -1,5 +1,7 @@ import asyncio, copy, importlib, os, time from datetime import datetime +from typing import Final +from unittest.mock import patch import pytest @@ -50,6 +52,49 @@ def test_log_success_event_counts_a_response_with_no_completion_tokens(): assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1} +@pytest.mark.parametrize( + ("candidate_model", "candidate_price", "expected_id"), + ( + ("openai/cheap", None, "candidate"), + ("other/cheap", None, "reference"), + (None, None, "reference"), + ("unknown", None, "reference"), + ("ollama/unknown", None, "reference"), + ("custom/cheap", {"input_cost_per_token": 0.01, "output_cost_per_token": 0.02}, "candidate"), + ("openai/cheap", {"input_cost_per_token": 2.0, "output_cost_per_token": 3.0}, "reference"), + ("custom/broken", {"input_cost_per_token": "invalid"}, "reference"), + ("custom/null", {"input_cost_per_token": None, "output_cost_per_token": None}, "reference"), + ), +) +@pytest.mark.asyncio +@pytest.mark.parametrize("reverse_order", [False, True], ids=["candidate-first", "reference-first"]) +async def test_provider_prefix_uses_matching_static_prices_and_preserves_exact_entries( + candidate_model: str | None, + candidate_price: dict[str, float | str | None] | None, + expected_id: str, + reverse_order: bool, +) -> None: + deployments: Final = [ + {"litellm_params": {"model": candidate_model}, "model_info": {"id": "candidate"}}, + {"litellm_params": {"model": "openai/reference"}, "model_info": {"id": "reference"}}, + ] + prices: Final = { + "cheap": {"input_cost_per_token": 0.01, "output_cost_per_token": 0.02, "litellm_provider": "openai"}, + "reference": {"input_cost_per_token": 0.2, "output_cost_per_token": 0.3, "litellm_provider": "openai"}, + **({candidate_model: candidate_price} if candidate_model is not None and candidate_price is not None else {}), + } + handler: Final = LowestCostLoggingHandler(router_cache=DualCache()) + + with patch.dict(litellm.model_cost, prices, clear=True), patch("litellm.get_model_info") as metadata_lookup: + selected: Final = await handler.async_get_available_deployments( + model_group="test-group", healthy_deployments=deployments[::-1] if reverse_order else deployments + ) + + expected: Final = next(deployment for deployment in deployments if deployment["model_info"]["id"] == expected_id) + assert selected is expected + metadata_lookup.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) async def test_log_success_event_keeps_cost_bookkeeping_out_of_the_latency_routing_entry(use_async: bool):