mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
fix: use deployment pricing for cost-based routing
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
joshua
parent
fb051f9647
commit
6eb6277fda
3 changed files with 174 additions and 21 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue