From 60c8b3dccb4589f206430d837921aa9816d8920b Mon Sep 17 00:00:00 2001 From: jagjeet-singh-23 Date: Sat, 5 Sep 2026 13:19:14 +0530 Subject: [PATCH] fix(budget-limiter): fall back to logging payload for custom_llm_provider `/v1/chat/completions` injects `custom_llm_provider` into `litellm_params`, but `/v1/messages` and `/v1/embeddings` do not. `RouterBudgetLimiting.async_log_success_event` read it only from `litellm_params` and raised `ValueError("custom_llm_provider is required")` before incrementing any spend counter, so provider budgets never accrued on those routes. Deployment-level `max_budget` sits below the same guard, so it was silently disabled too. The standard logging payload already carries `custom_llm_provider` for every route -- it is built from the top-level call details rather than from `litellm_params` -- so fall back to it. `litellm_params` keeps precedence, and the `ValueError` is preserved when neither source has a provider. Fixes #26701 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_015ftJdiLtpTbjoZaNuqX6dq --- litellm/router_strategy/budget_limiter.py | 10 +- ...test_budget_limiter_provider_derivation.py | 125 ++++++++++++++++++ 2 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter_provider_derivation.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index a8d51f95e45..25c5be0e5d3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -408,7 +408,15 @@ 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) + # `/v1/chat/completions` injects `custom_llm_provider` into `litellm_params`, + # but `/v1/messages` and `/v1/embeddings` do not. The standard logging payload + # carries it for every route (it is built from the top-level call details), so + # fall back to that instead of raising before any spend counter is incremented. + # See https://github.com/BerriAI/litellm/issues/26701. + custom_llm_provider: Final[str | None] = ( + kwargs.get("litellm_params", {}).get("custom_llm_provider", None) + or standard_logging_payload["custom_llm_provider"] + ) if custom_llm_provider is None: raise ValueError("custom_llm_provider is required") diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_provider_derivation.py b/tests/test_litellm/router_strategy/test_budget_limiter_provider_derivation.py new file mode 100644 index 00000000000..4177f655c64 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter_provider_derivation.py @@ -0,0 +1,125 @@ +""" +Regression tests for #26701. + +`/v1/chat/completions` injects `custom_llm_provider` into `litellm_params`, but +`/v1/messages` and `/v1/embeddings` do not. Before the fix, +`RouterBudgetLimiting.async_log_success_event` raised +`ValueError("custom_llm_provider is required")` before incrementing any spend +counter, so provider budgets (and deployment-level `max_budget`, which sits +below the same guard) silently never accrued on those routes. + +The standard logging payload carries `custom_llm_provider` for every route -- +it is built from the top-level call details rather than from `litellm_params` -- +so the limiter falls back to it. +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.caching.caching import DualCache +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + + +def _make_limiter() -> RouterBudgetLimiting: + return RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={ + "anthropic": {"budget_limit": 5.0, "time_period": "24h"}, + "openai": {"budget_limit": 5.0, "time_period": "24h"}, + }, + ) + + +def _payload(provider: str | None) -> dict: + return { + "response_cost": 0.001, + "model_id": "deployment-abc123", + "custom_llm_provider": provider, + } + + +def _patched(limiter: RouterBudgetLimiting): + """Isolate the provider-budget path from deployment/tag budget lookups.""" + return ( + patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock), + patch.object(limiter, "_get_budget_config_for_deployment", return_value=None), + patch.object(limiter, "_get_budget_config_for_tag", return_value=None), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "model", "provider"), + [ + ("/v1/messages", "anthropic/claude-haiku-4-5-20251001", "anthropic"), + ("/v1/embeddings", "openai/text-embedding-3-small", "openai"), + ], +) +async def test_provider_falls_back_to_logging_payload(route: str, model: str, provider: str): + """ + Routes that omit `custom_llm_provider` from `litellm_params` must still + accrue spend, using the provider recorded on the standard logging payload. + """ + limiter = _make_limiter() + kwargs = { + "model": model, + # `litellm_params` on these routes carries no `custom_llm_provider`. + "litellm_params": {"model": model}, + "standard_logging_object": _payload(provider), + } + + increment, _dep, _tag = _patched(limiter) + with increment as mock_increment, _dep, _tag: + # Before the fix this raised ValueError before any counter moved. + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + + assert mock_increment.await_count >= 1, f"{route}: provider spend was never incremented" + spend_keys = [call.kwargs.get("spend_key", "") for call in mock_increment.await_args_list] + assert any(f"provider_spend:{provider}:" in key for key in spend_keys), ( + f"{route}: expected a provider_spend key for {provider!r}, got {spend_keys}" + ) + + +@pytest.mark.asyncio +async def test_litellm_params_provider_takes_precedence(): + """ + `/v1/chat/completions` already sets `custom_llm_provider` in `litellm_params`; + that value must keep winning over the payload so behaviour is unchanged. + """ + limiter = _make_limiter() + kwargs = { + "model": "claude-haiku-direct", + "litellm_params": { + "model": "anthropic/claude-haiku-4-5-20251001", + "custom_llm_provider": "anthropic", + }, + # Deliberately disagrees, to prove precedence rather than coincidence. + "standard_logging_object": _payload("openai"), + } + + increment, _dep, _tag = _patched(limiter) + with increment as mock_increment, _dep, _tag: + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) + + spend_keys = [call.kwargs.get("spend_key", "") for call in mock_increment.await_args_list] + assert any("provider_spend:anthropic:" in key for key in spend_keys), spend_keys + assert not any("provider_spend:openai:" in key for key in spend_keys), spend_keys + + +@pytest.mark.asyncio +async def test_missing_provider_everywhere_still_raises(): + """ + With no provider in `litellm_params` and none on the payload there is + nothing to attribute spend to; the existing contract must be preserved. + """ + limiter = _make_limiter() + kwargs = { + "model": "totally-unresolvable-garbage-xyz-123", + "litellm_params": {}, + "standard_logging_object": _payload(None), + } + + increment, _dep, _tag = _patched(limiter) + with increment, _dep, _tag, pytest.raises(ValueError, match="custom_llm_provider is required"): + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None)