scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge 242e96c0ee into a6816f0e96

Browse source
This commit is contained in:
Jagjeet Singh 2026-08-27 17:33:01 -05:00 • committed by GitHub
commit a2d8e4d717
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 215 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/router_strategy/budget_limiter.py
View file

@ -408,6 +408,37 @@ 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:
# /v1/messages and /v1/embeddings routes do not inject
# custom_llm_provider into litellm_params the way
# /v1/chat/completions does. Derive it from the model string
# so budget tracking is not silently skipped for those routes.
# See: https://github.com/BerriAI/litellm/issues/26701
_litellm_params = kwargs.get("litellm_params") or {}
_model = (
_litellm_params.get("model", "")
if isinstance(_litellm_params, dict)
else getattr(_litellm_params, "model", "") or ""
)
# litellm_params may not carry `model` for /v1/messages and
# /v1/embeddings routes (LoggedLiteLLMParams gap). Fall back to
# the top-level model field and then the standard logging payload
# so budget enforcement is not silently skipped.
# See: https://github.com/BerriAI/litellm/issues/26701
if not _model:
_model = kwargs.get("model", "") or standard_logging_payload.get("model", "") or ""
try:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=str(_model),
litellm_params=_LiteLLMParamsDictView(_litellm_params if isinstance(_litellm_params, dict) else {}),
)
except Exception as e: # noqa: BLE001
verbose_router_logger.debug(
"RouterBudgetLimiting: could not derive custom_llm_provider from model string %r: %s",
_model,
e,
)
custom_llm_provider = None
if custom_llm_provider is None:
raise ValueError("custom_llm_provider is required")

184
tests/test_litellm/router_strategy/test_budget_limiter_provider_derivation.py Normal file
View file

@ -0,0 +1,184 @@
"""
Regression tests for #26701
/v1/messages and /v1/embeddings routes don't inject custom_llm_provider into
litellm_params (unlike /v1/chat/completions). async_log_success_event must
derive the provider from the model string instead of raising ValueError.
"""
import pytest
from unittest.mock import AsyncMock, patch
from litellm.caching.caching import DualCache
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
def _make_limiter():
return RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={
"anthropic": {"budget_limit": 5.0, "time_period": "24h"}
},
)
def _kwargs_without_custom_provider(model_str: str) -> dict:
"""Simulates what /v1/messages and /v1/embeddings produce — no custom_llm_provider."""
return {
"model": model_str,
"litellm_params": {
"model": model_str,
# custom_llm_provider deliberately absent — this is the bug scenario
},
"standard_logging_object": {
"response_cost": 0.001,
"model_id": "deployment-abc123",
},
}
@pytest.mark.asyncio
async def test_messages_route_no_custom_provider_does_not_raise():
"""
/v1/messages kwargs missing custom_llm_provider must not raise ValueError.
Provider should be derived from the model string 'anthropic/...' instead.
"""
limiter = _make_limiter()
kwargs = _kwargs_without_custom_provider("anthropic/claude-haiku-4-5-20251001")
with (
patch.object(
limiter, "_increment_spend_for_key", new_callable=AsyncMock
) as mock_increment,
patch.object(limiter, "_get_budget_config_for_deployment", return_value=None),
patch.object(limiter, "_get_budget_config_for_tag", return_value=None),
):
# Before the fix this raises: ValueError("custom_llm_provider is required")
await limiter.async_log_success_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
# Provider was correctly derived → spend increment was called
mock_increment.assert_called_once()
@pytest.mark.asyncio
async def test_embeddings_route_no_custom_provider_does_not_raise():
"""/v1/embeddings has the same gap — openai provider should be derived cleanly."""
limiter = _make_limiter()
kwargs = _kwargs_without_custom_provider("openai/text-embedding-3-small")
with (
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),
):
await limiter.async_log_success_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
@pytest.mark.asyncio
async def test_chat_completions_with_explicit_provider_still_works():
"""
/v1/chat/completions path that already sets custom_llm_provider must
continue to work exactly as before — no regression.
"""
limiter = _make_limiter()
kwargs = {
"model": "claude-haiku-direct",
"litellm_params": {
"model": "anthropic/claude-haiku-4-5-20251001",
"custom_llm_provider": "anthropic", # set by chat/completions route
},
"standard_logging_object": {
"response_cost": 0.001,
"model_id": "deployment-abc123",
},
}
with (
patch.object(
limiter, "_increment_spend_for_key", new_callable=AsyncMock
) as mock_increment,
patch.object(limiter, "_get_budget_config_for_deployment", return_value=None),
patch.object(limiter, "_get_budget_config_for_tag", return_value=None),
):
await limiter.async_log_success_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
mock_increment.assert_called_once()
@pytest.mark.asyncio
async def test_unresolvable_model_string_raises_value_error_with_debug_log():
"""
When both the dict lookup and litellm.get_llm_provider() fail
(completely unrecognisable model string), ValueError must still be raised
and the debug log must fire — covering the except branch added in #26701.
"""
limiter = _make_limiter()
kwargs = _kwargs_without_custom_provider("totally-unresolvable-garbage-xyz-123")
with (
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),
):
with 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,
)
@pytest.mark.asyncio
async def test_provider_derived_from_top_level_model_when_litellm_params_missing_model():
"""
Regression for veria-ai review comment on #32180:
LoggedLiteLLMParams for /v1/messages does not include `model`, so
litellm_params.get("model") returns "". Must fall back to kwargs["model"]
so that provider spend is actually incremented (budget enforced).
"""
limiter = _make_limiter()
kwargs = {
"model": "anthropic/claude-haiku-4-5-20251001", # top-level, always present
"litellm_params": {
# model key absent — this is the LoggedLiteLLMParams gap
},
"standard_logging_object": {
"response_cost": 0.001,
"model_id": "deployment-abc123",
},
}
with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment, \
patch.object(limiter, "_get_budget_config_for_deployment", return_value=None), \
patch.object(limiter, "_get_budget_config_for_tag", return_value=None):
await limiter.async_log_success_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
# Budget spend MUST be incremented — if it isn't, budget enforcement is broken
mock_increment.assert_called_once()