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

@@ -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}"

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