diff --git a/litellm/router.py b/litellm/router.py index cdede46ef5c..f95cb60df94 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1508,9 +1508,23 @@ class Router: routing_args=routing_strategy_args, ) case RoutingStrategy.COST_BASED.value: + def price_lookup( + model_id: str, model_name: str + ) -> tuple[float | None, float | None]: + model_info: Final = self.get_deployment_model_info( + model_id=model_id, model_name=model_name + ) + if model_info is None: + return None, None + return ( + model_info.get("input_cost_per_token"), + model_info.get("output_cost_per_token"), + ) + selector = LowestCostLoggingHandler( router_cache=self.cache, routing_args={}, + price_lookup=price_lookup, ) case _: pass diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index d567b6acccc..00daa71fffa 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -1,7 +1,8 @@ #### What this does #### # picks based on response time (for streaming, this is time to first token) +from collections.abc import Callable from datetime import datetime -from typing import Final +from typing import Final, cast # noqa: TID251 # legacy deployment payloads have no shared typed boundary import litellm from litellm import ModelResponse, token_counter, verbose_logger @@ -12,13 +13,46 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +def _deployment_cost( + model_id: str, + model_name: str, + input_price_override: float | None, + output_price_override: float | None, + price_lookup: Callable[[str, str], tuple[float | None, float | None]], +) -> float: + lookup_input_price, lookup_output_price = price_lookup(model_id, model_name) + input_price: Final = ( + input_price_override if input_price_override is not None else lookup_input_price + ) + output_price: Final = ( + output_price_override if output_price_override is not None else lookup_output_price + ) + input_cost: Final = input_price if input_price is not None else 5.0 + output_cost: Final = output_price if output_price is not None else 5.0 + return input_cost + output_cost + + +def _default_price_lookup(model_id: str, model_name: str) -> tuple[float | None, float | None]: + model_cost: Final = litellm.model_cost.get(model_id) or litellm.model_cost.get(model_name) + if model_cost is None: + return None, None + return model_cost.get("input_cost_per_token"), model_cost.get("output_cost_per_token") + + class LowestCostLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 logged_failure: int = 0 + price_lookup: Callable[[str, str], tuple[float | None, float | None]] - def __init__(self, router_cache: DualCache, routing_args: dict = {}): + def __init__( + self, + router_cache: DualCache, + routing_args: dict = {}, + price_lookup: Callable[[str, str], tuple[float | None, float | None]] = _default_price_lookup, + ): self.router_cache = router_cache + self.price_lookup = price_lookup @with_service_target("router_usage") def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -243,28 +277,24 @@ class LowestCostLoggingHandler(CustomLogger): or _deployment.get("model_info", {}).get("rpm", None) or float("inf") ) - item_litellm_model_name = _deployment.get("litellm_params", {}).get("model") - item_litellm_model_cost_map = litellm.model_cost.get(item_litellm_model_name, {}) - - # check if user provided input_cost_per_token and output_cost_per_token in litellm_params - item_input_cost = None - item_output_cost = None - if _deployment.get("litellm_params", {}).get("input_cost_per_token", None): - item_input_cost = _deployment.get("litellm_params", {}).get("input_cost_per_token") - - if _deployment.get("litellm_params", {}).get("output_cost_per_token", None): - item_output_cost = _deployment.get("litellm_params", {}).get("output_cost_per_token") - - if item_input_cost is None: - item_input_cost = item_litellm_model_cost_map.get("input_cost_per_token", 5.0) - - if item_output_cost is None: - item_output_cost = item_litellm_model_cost_map.get("output_cost_per_token", 5.0) + item_cost = _deployment_cost( + model_id=cast(str, item), # cast-ok: router cache keys come from deployment ids + model_name=cast( # cast-ok: router validates deployment model names + str, _deployment.get("litellm_params", {}).get("model") or "" + ), + input_price_override=cast( # cast-ok: router validates configured token prices + float | None, + _deployment.get("litellm_params", {}).get("input_cost_per_token"), + ), + output_price_override=cast( # cast-ok: router validates configured token prices + float | None, + _deployment.get("litellm_params", {}).get("output_cost_per_token"), + ), + price_lookup=self.price_lookup, + ) # if litellm["model"] is not in model_cost map -> use item_cost = $10 - item_cost = item_input_cost + item_output_cost - item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) diff --git a/tests/unit/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py index b53ec6e6ba2..0edee190a9c 100644 --- a/tests/unit/router_strategy/test_lowest_cost.py +++ b/tests/unit/router_strategy/test_lowest_cost.py @@ -1,12 +1,16 @@ import asyncio, copy, importlib, os, time from datetime import datetime +from typing import Final, Literal import pytest import litellm from litellm.caching.caching import DualCache +from litellm.router import Router +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -357,3 +361,108 @@ async def test_get_available_endpoints_tpm_rpm_check_async(ans_rpm): assert (d_ans and d_ans["model_info"]["id"]) == ans print("selected deployment:", d_ans) + + +def _cost_routing_deployment( + deployment_id: str, + model: str, + pricing_source: Literal["litellm_params", "model_info"] | None = None, + prices: tuple[float, float] | None = None, + base_model: str | None = None, +) -> DeploymentTypedDict: + pricing: Final = ( + {"input_cost_per_token": prices[0], "output_cost_per_token": prices[1]} + if prices is not None + else {} + ) + litellm_params: Final[LiteLLMParamsTypedDict] = ( + { + "model": model, + "mock_response": "ok", + "input_cost_per_token": prices[0], + "output_cost_per_token": prices[1], + } + if pricing_source == "litellm_params" and prices is not None + else {"model": model, "mock_response": "ok"} + ) + return { + "model_name": "cost-pricing", + "litellm_params": litellm_params, + "model_info": { + "id": deployment_id, + **({"base_model": base_model} if base_model is not None else {}), + **(pricing if pricing_source == "model_info" else {}), + }, + } + + +def _cost_routing_deployments( + pricing_case: Literal["litellm_params", "model_info", "base_model", "zero_price"], +) -> list[DeploymentTypedDict]: + match pricing_case: + case "litellm_params": + return [ + _cost_routing_deployment( + "litellm-params-expensive", "openai/gpt-4o-mini", "litellm_params", (3e-6, 3e-6) + ), + _cost_routing_deployment( + "litellm-params-cheap", "openai/gpt-4o-mini", "litellm_params", (4e-7, 4e-7) + ), + ] + case "model_info": + return [ + _cost_routing_deployment( + "model-info-expensive", "openai/gpt-4o-mini", "model_info", (3e-6, 3e-6) + ), + _cost_routing_deployment( + "model-info-cheap", "openai/gpt-4o-mini", "model_info", (4e-7, 4e-7) + ), + ] + case "base_model": + return [ + _cost_routing_deployment( + "base-model-expensive", "azure/my-4o", base_model="azure/gpt-4o" + ), + _cost_routing_deployment( + "base-model-cheap", "azure/my-mini", base_model="azure/gpt-4o-mini" + ), + ] + case "zero_price": + return [ + _cost_routing_deployment( + "nonzero-price", "openai/gpt-4o-mini", "model_info", (1e-7, 1e-7) + ), + _cost_routing_deployment( + "zero-price", "openai/gpt-4o-mini", "model_info", (0.0, 0.0) + ), + ] + + +@pytest.mark.usefixtures("isolate_litellm_state", "setup_and_teardown") +@pytest.mark.asyncio +@pytest.mark.parametrize( + "pricing_case", + ("litellm_params", "model_info", "base_model", "zero_price"), +) +async def test_cost_based_router_uses_deployment_model_info_pricing( + pricing_case: Literal["litellm_params", "model_info", "base_model", "zero_price"], +) -> None: + model_list: Final = _cost_routing_deployments(pricing_case) + cheap_deployment_id: Final = { + "litellm_params": "litellm-params-cheap", + "model_info": "model-info-cheap", + "base_model": "base-model-cheap", + "zero_price": "zero-price", + }[pricing_case] + router: Final = Router( + model_list=model_list, + routing_strategy="cost-based-routing", + num_retries=0, + ) + + response: Final = await router.acompletion( + model="cost-pricing", + messages=[{"role": "user", "content": "choose by deployment price"}], + ) + + assert get_hidden_params_dict(response)["model_id"] == cheap_deployment_id