Merge 6d432966f4 into f2118fc767

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

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):