From 8327cd6d4724b6218f85562ba164f0cb2cb71ef6 Mon Sep 17 00:00:00 2001 From: daqiangganjun <93830914+daqiangganjun@users.noreply.github.com> Date: Sat, 26 Sep 2026 05:06:49 +0800 Subject: [PATCH] fix(router): count provider budget spend on every API surface (#38172) * fix(router): count provider budget spend on every API surface RouterBudgetLimiting read custom_llm_provider from litellm_params, which only chat completions populates. Responses, anthropic_messages, embedding and rerank calls raised inside the success callback before any spend was recorded, so those budgets never moved and a ceiling made up mostly of that traffic was never hit. Read the provider from the standard logging payload, which every surface fills in. Dropping the raise also stops one missing field from taking the deployment and tag budgets down with it. * chore(router): drop the inline comment and type the budget limiter test helper --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router_strategy/budget_limiter.py | 10 +- .../router_strategy/test_budget_limiter.py | 137 ++++++++++++++++++ 2 files changed, 142 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 4e84bded9de..64252cbbfb3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -502,12 +502,12 @@ class RouterBudgetLimiting(CustomLogger): response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) - custom_llm_provider: Final[str] = kwargs.get("litellm_params", {}).get("custom_llm_provider", None) - if custom_llm_provider is None: - raise ValueError("custom_llm_provider is required") + custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") - budget_config: Final = self._get_budget_config_for_provider(custom_llm_provider) - if budget_config: + budget_config: Final = ( + self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None + ) + if custom_llm_provider is not None and budget_config is not None: # increment spend for provider spend_key: Final = f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}" start_time_key: Final = f"provider_budget_start_time:{custom_llm_provider}" diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..62de1586fdd --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,137 @@ +""" +Spend tracking in RouterBudgetLimiting.async_log_success_event. + +Only chat completions puts custom_llm_provider into litellm_params. The responses, +anthropic_messages, embedding and rerank surfaces leave it unset, which used to make +the callback raise before any spend was recorded, so those budgets never moved. +""" + +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + + +@pytest.fixture +def disable_budget_sync(monkeypatch): + async def noop(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + noop, + ) + + +def _success_kwargs( + *, + provider_in_litellm_params: str | None, + provider_in_payload: str | None, + call_type: str = "aresponses", + response_cost: float = 0.25, + model_id: str = "deployment-1", +) -> dict[str, object]: + provider_params: Final[dict[str, str]] = ( + {} if provider_in_litellm_params is None else {"custom_llm_provider": provider_in_litellm_params} + ) + litellm_params: Final[dict[str, str]] = {"model": "openai/gpt-4o", **provider_params} + + return { + "call_type": call_type, + "litellm_params": litellm_params, + "standard_logging_object": { + "response_cost": response_cost, + "model_id": model_id, + "custom_llm_provider": provider_in_payload, + }, + } + + +async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object]) -> None: + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"]) +async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type): + """Non-chat surfaces carry the provider only on the standard logging payload.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs( + provider_in_litellm_params=None, + provider_in_payload="openai", + call_type=call_type, + ), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.25 + + +@pytest.mark.asyncio +async def test_chat_completions_spend_still_tracked(disable_budget_sync): + """Chat completions fills in both sources and must keep accumulating.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs( + provider_in_litellm_params="openai", + provider_in_payload="openai", + call_type="acompletion", + ), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == 0.25 + + +@pytest.mark.asyncio +async def test_budget_of_other_provider_is_untouched(disable_budget_sync): + """A provider without its own budget must not bleed into a configured one.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + + await _log_success( + limiter, + _success_kwargs(provider_in_litellm_params=None, provider_in_payload="anthropic"), + ) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") in (None, 0.0) + + +@pytest.mark.asyncio +async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_budget_sync): + """An unresolvable provider must not abort the deployment and tag budgets that follow it.""" + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config=None, + model_list=[ + { + "model_name": "some-model", + "litellm_params": { + "model": "openai/gpt-4o", + "max_budget": 10.0, + "budget_duration": "1d", + }, + "model_info": {"id": "deployment-1"}, + } + ], + ) + + await _log_success( + limiter, + _success_kwargs(provider_in_litellm_params=None, provider_in_payload=None), + ) + + assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25