diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 631b0c3df3d..2c5cc6835e8 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -104,6 +104,10 @@ def _sum_increments_by_key(operations: Sequence[RedisPipelineIncrementOperation] return MappingProxyType({key: sum(operation["increment_value"] for operation in group) for key, group in by_key}) +def _is_over_budget(spend: float, max_budget: float | None) -> bool: + return max_budget is not None and spend >= max_budget + + class RouterBudgetLimiting(CustomLogger): def __init__( self, @@ -239,16 +243,15 @@ class RouterBudgetLimiting(CustomLogger): provider = self._get_llm_provider_for_deployment(deployment) if provider in provider_configs: config = provider_configs[provider] - if config.max_budget is None: - continue current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0) - self._track_provider_remaining_budget_prometheus( - provider=provider, - spend=current_spend, - budget_limit=config.max_budget, - ) + if config.max_budget is not None: + self._track_provider_remaining_budget_prometheus( + provider=provider, + spend=current_spend, + budget_limit=config.max_budget, + ) - if config.max_budget and current_spend >= config.max_budget: + if _is_over_budget(current_spend, config.max_budget): debug_msg = f"Exceeded budget for provider {provider}: {current_spend} >= {config.max_budget}" deployment_above_budget_info += f"{debug_msg}\n" is_within_budget = False @@ -263,8 +266,8 @@ class RouterBudgetLimiting(CustomLogger): if model_id in deployment_configs: config = deployment_configs[model_id] current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0) - if config.max_budget and current_spend >= config.max_budget: - debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.budget_duration}" + if _is_over_budget(current_spend, config.max_budget): + debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.max_budget}" verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" is_within_budget = False @@ -278,7 +281,7 @@ class RouterBudgetLimiting(CustomLogger): f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}", 0.0, ) - if _tag_budget_config.max_budget and _tag_spend >= _tag_budget_config.max_budget: + if _is_over_budget(_tag_spend, _tag_budget_config.max_budget): debug_msg = f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, tag_budget_limit={_tag_budget_config.max_budget}" verbose_router_logger.debug(debug_msg) deployment_above_budget_info += f"{debug_msg}\n" diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 62de1586fdd..92c6163c864 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -10,6 +10,7 @@ from typing import Final import pytest +from litellm import Router from litellm.caching.caching import DualCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -135,3 +136,122 @@ async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_b ) assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 + + +_OPENAI_DEPLOYMENT: Final = { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"id": "openai-dep"}, +} +_ANTHROPIC_DEPLOYMENT: Final = { + "model_name": "chat", + "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "sk-fake"}, + "model_info": {"id": "anthropic-dep"}, +} + + +def _openai_deployment_with_budget(cap: float | None) -> dict[str, object]: + budget: Final = {"budget_duration": "1d"} if cap is None else {"max_budget": cap, "budget_duration": "1d"} + return {**_OPENAI_DEPLOYMENT, "litellm_params": {**_OPENAI_DEPLOYMENT["litellm_params"], **budget}} + + +def _limiter_capping_openai(scope: str, cap: float | None, monkeypatch: pytest.MonkeyPatch) -> RouterBudgetLimiting: + budget: Final = {"time_period": "1d"} if cap is None else {"budget_limit": cap, "time_period": "1d"} + if scope == "provider": + return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={"openai": budget}) + if scope == "deployment": + return RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config=None, + model_list=[_openai_deployment_with_budget(cap), _ANTHROPIC_DEPLOYMENT], + ) + monkeypatch.setattr("litellm.tag_budget_config", {"prod": budget}) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None) + + +async def _record_openai_spend(limiter: RouterBudgetLimiting, cost: float) -> None: + await _log_success( + limiter, + { + "call_type": "acompletion", + "metadata": {"tags": ["prod"]}, + "litellm_params": {"model": "openai/gpt-4o", "custom_llm_provider": "openai"}, + "standard_logging_object": { + "response_cost": cost, + "model_id": "openai-dep", + "custom_llm_provider": "openai", + }, + }, + ) + + +async def _routable_ids(limiter: RouterBudgetLimiting, deployments: list[dict[str, object]]) -> list[str]: + kept: Final = await limiter.async_filter_deployments( + model="chat", + healthy_deployments=deployments, + messages=None, + request_kwargs={"metadata": {"tags": ["prod"]}}, + ) + return [deployment["model_info"]["id"] for deployment in kept] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["provider", "deployment", "tag"]) +@pytest.mark.parametrize( + ("cap", "spend", "openai_routable"), + [ + (0.0, 0.0, False), + (0.01, 0.0, True), + (0.01, 0.009, True), + (0.01, 0.01, False), + (None, 50.0, True), + ], +) +async def test_zero_cap_blocks_and_only_an_unset_cap_is_unlimited( + disable_budget_sync: None, + monkeypatch: pytest.MonkeyPatch, + scope: str, + cap: float | None, + spend: float, + openai_routable: bool, +) -> None: + limiter: Final = _limiter_capping_openai(scope, cap, monkeypatch) + if spend > 0: + await _record_openai_spend(limiter, spend) + openai: Final = _openai_deployment_with_budget(cap) if scope == "deployment" else _OPENAI_DEPLOYMENT + + if scope == "tag" and not openai_routable: + with pytest.raises(ValueError, match=r"Exceeded budget for tag='prod'"): + await _routable_ids(limiter, [openai, _ANTHROPIC_DEPLOYMENT]) + return + + routable: Final = await _routable_ids(limiter, [openai, _ANTHROPIC_DEPLOYMENT]) + + expected: Final = ["openai-dep", "anthropic-dep"] if openai_routable else ["anthropic-dep"] + assert routable == expected, f"{scope} cap={cap} spend={spend}" + + +@pytest.mark.asyncio +async def test_router_rejects_requests_to_a_provider_capped_at_zero(disable_budget_sync: None) -> None: + router: Final = Router( + model_list=[_OPENAI_DEPLOYMENT], + provider_budget_config={"openai": {"budget_limit": 0, "time_period": "1d"}}, + ) + + with pytest.raises(ValueError, match=r"Exceeded budget for provider openai: 0\.0 >= 0\.0"): + await router.acompletion(model="chat", messages=[{"role": "user", "content": "hi"}], mock_response="served") + + +@pytest.mark.asyncio +async def test_router_serves_a_provider_with_a_period_but_no_cap(disable_budget_sync: None) -> None: + router: Final = Router( + model_list=[_OPENAI_DEPLOYMENT], + provider_budget_config={"openai": {"time_period": "1d"}}, + ) + + served: Final = await router.acompletion( + model="chat", messages=[{"role": "user", "content": "hi"}], mock_response="served" + ) + + assert served.choices[0].message.content == "served"