Merge 988ebe5bce into 78ff5ac9cd

litellm/router_strategy/budget_limiter.py

@@ -408,12 +408,14 @@ 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")
+        # litellm_params only carries the provider on chat completions; the payload carries it on every
+        # surface, so read it from there instead of skipping the budget for responses, messages and embeddings.
+        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,134 @@
+"""
+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.
+"""
+
+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:
+    litellm_params = {"model": "openai/gpt-4o"}
+    if provider_in_litellm_params is not None:
+        litellm_params["custom_llm_provider"] = provider_in_litellm_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) -> 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