mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy/budgets): reject requests at the exact hard-budget limit
This commit is contained in:
parent
65ca095d4d
commit
ede4c63547
5 changed files with 178 additions and 6 deletions
|
|
@ -1088,7 +1088,7 @@ async def _check_end_user_budget(
|
|||
max_budget=end_user_budget,
|
||||
fallback_authoritative=True,
|
||||
)
|
||||
if end_user_spend > end_user_budget:
|
||||
if end_user_spend >= end_user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=end_user_spend,
|
||||
max_budget=end_user_budget,
|
||||
|
|
@ -3900,7 +3900,7 @@ async def _team_max_budget_check(
|
|||
max_budget=team_object.max_budget,
|
||||
)
|
||||
|
||||
if math.isfinite(team_object.max_budget) and spend > team_object.max_budget:
|
||||
if math.isfinite(team_object.max_budget) and spend >= team_object.max_budget:
|
||||
if valid_token:
|
||||
call_info = CallInfo(
|
||||
token=valid_token.token,
|
||||
|
|
@ -4320,7 +4320,7 @@ async def _tag_max_budget_check(
|
|||
max_budget=tag_object.litellm_budget_table.max_budget,
|
||||
fallback_authoritative=True,
|
||||
)
|
||||
if tag_spend <= tag_object.litellm_budget_table.max_budget:
|
||||
if tag_spend < tag_object.litellm_budget_table.max_budget:
|
||||
continue
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=tag_spend,
|
||||
|
|
|
|||
|
|
@ -1793,7 +1793,7 @@ async def _user_api_key_auth_builder(
|
|||
fallback_spend=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
)
|
||||
if team_member_spend > team_member_budget:
|
||||
if team_member_spend >= team_member_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
if (
|
||||
_current_spend is not None
|
||||
and _current_model_budget_info.max_budget is not None
|
||||
and _current_spend > _current_model_budget_info.max_budget
|
||||
and _current_spend >= _current_model_budget_info.max_budget
|
||||
):
|
||||
raise litellm.BudgetExceededError(
|
||||
message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
|
||||
|
|
@ -134,7 +134,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
if (
|
||||
_current_spend is not None
|
||||
and _current_model_budget_info.max_budget is not None
|
||||
and _current_spend > _current_model_budget_info.max_budget
|
||||
and _current_spend >= _current_model_budget_info.max_budget
|
||||
):
|
||||
raise litellm.BudgetExceededError(
|
||||
message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}",
|
||||
|
|
|
|||
|
|
@ -2978,6 +2978,94 @@ async def test_team_member_budget_check_reads_from_spend_counter():
|
|||
assert exc_info.value.current_cost == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_budget_check_rejects_at_exact_limit():
|
||||
"""Team budget must be rejected when spend equals max_budget (hard limit)."""
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
team_object = LiteLLM_TeamTable(team_id="test-team", spend=0.0, max_budget=1.0)
|
||||
valid_token = UserAPIKeyAuth(token="test-token", team_id="test-team")
|
||||
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
|
||||
proxy_logging_obj.budget_alerts = AsyncMock()
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
if counter_key == "spend:team:test-team":
|
||||
return 1.0
|
||||
return fallback_spend
|
||||
|
||||
with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _team_max_budget_check(
|
||||
team_object=team_object,
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
assert exc_info.value.current_cost == 1.0
|
||||
assert exc_info.value.max_budget == 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_budget_check_rejects_at_exact_limit():
|
||||
"""End-user budget must be rejected when spend equals max_budget (hard limit)."""
|
||||
end_user_object = LiteLLM_EndUserTable(
|
||||
user_id="customer-1",
|
||||
blocked=False,
|
||||
spend=0.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
)
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
if counter_key == "spend:end_user:customer-1":
|
||||
return 1.0
|
||||
return fallback_spend
|
||||
|
||||
with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_end_user_budget(
|
||||
end_user_obj=end_user_object,
|
||||
route="/chat/completions",
|
||||
)
|
||||
assert exc_info.value.current_cost == 1.0
|
||||
assert exc_info.value.max_budget == 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_budget_check_rejects_at_exact_limit():
|
||||
"""Tag budget must be rejected when spend equals max_budget (hard limit)."""
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
tag_object = LiteLLM_TagTable(
|
||||
tag_name="paid-tag",
|
||||
spend=0.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
)
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
if counter_key == "spend:tag:paid-tag":
|
||||
return 1.0
|
||||
return fallback_spend
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_checks.get_tag_objects_batch",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"paid-tag": tag_object},
|
||||
),
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _tag_max_budget_check(
|
||||
request_body={"metadata": {"tags": ["paid-tag"]}},
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
|
||||
valid_token=UserAPIKeyAuth(token="test-token"),
|
||||
)
|
||||
assert exc_info.value.current_cost == 1.0
|
||||
assert exc_info.value.max_budget == 1.0
|
||||
|
||||
|
||||
class TestGuardrailModificationCheck:
|
||||
"""Defense-in-depth: `_guardrail_modification_check` must 403 when the
|
||||
caller's metadata attempts to modify any guardrail-related key and the
|
||||
|
|
|
|||
|
|
@ -0,0 +1,84 @@
|
|||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import (
|
||||
_PROXY_VirtualKeyModelMaxBudgetLimiter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_virtual_key_model_budget_rejects_at_exact_limit():
|
||||
"""Per-model virtual-key budget must reject when spend equals max_budget."""
|
||||
limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
limiter._get_virtual_key_spend_for_model = AsyncMock(return_value=10.0)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
token="test-token",
|
||||
model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}},
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await limiter.is_key_within_model_budget(
|
||||
user_api_key_dict=user_api_key_dict, model="gpt-4o"
|
||||
)
|
||||
assert exc_info.value.current_cost == 10.0
|
||||
assert exc_info.value.max_budget == 10.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_virtual_key_model_budget_admits_below_limit():
|
||||
"""Per-model virtual-key budget must admit when spend is below max_budget."""
|
||||
limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
limiter._get_virtual_key_spend_for_model = AsyncMock(return_value=9.99)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
token="test-token",
|
||||
model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}},
|
||||
)
|
||||
|
||||
assert (
|
||||
await limiter.is_key_within_model_budget(
|
||||
user_api_key_dict=user_api_key_dict, model="gpt-4o"
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_model_budget_rejects_at_exact_limit():
|
||||
"""Per-model end-user budget must reject when spend equals max_budget."""
|
||||
limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
limiter._get_end_user_spend_for_model = AsyncMock(return_value=10.0)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await limiter.is_end_user_within_model_budget(
|
||||
end_user_id="customer-1",
|
||||
end_user_model_max_budget={
|
||||
"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}
|
||||
},
|
||||
model="gpt-4o",
|
||||
)
|
||||
assert exc_info.value.current_cost == 10.0
|
||||
assert exc_info.value.max_budget == 10.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_model_budget_admits_below_limit():
|
||||
"""Per-model end-user budget must admit when spend is below max_budget."""
|
||||
limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
|
||||
limiter._get_end_user_spend_for_model = AsyncMock(return_value=9.99)
|
||||
|
||||
assert (
|
||||
await limiter.is_end_user_within_model_budget(
|
||||
end_user_id="customer-1",
|
||||
end_user_model_max_budget={
|
||||
"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}
|
||||
},
|
||||
model="gpt-4o",
|
||||
)
|
||||
is True
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue