scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix: add missing prisma_client guard to _check_model_max_budget for consistency with main budget-check block

Browse source
This commit is contained in:
Jason Kidd 2026-04-08 12:51:47 -07:00
parent 9df41015a4
commit 3503694ef1
No known key found for this signature in database
GPG key ID: 72BF942827539044
2 changed files with 24 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/auth/user_api_key_auth.py
View file

@ -1084,6 +1084,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
request_data=request_data,
route=route,
model_max_budget_limiter=model_max_budget_limiter,
prisma_client=prisma_client,
)
return valid_token
@ -1170,6 +1171,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
request_data=request_data,
route=route,
model_max_budget_limiter=model_max_budget_limiter,
prisma_client=prisma_client,
)
return _user_api_key_obj
@ -1806,6 +1808,7 @@ async def _check_model_max_budget(
request_data: dict,
route: str,
model_max_budget_limiter: "litellm.proxy.hooks.model_max_budget_limiter._PROXY_VirtualKeyModelMaxBudgetLimiter",
prisma_client: Optional["PrismaClient"] = None,
) -> None:
"""
Run per-model budget checks for both key-level and end-user-level budgets.
@ -1826,6 +1829,7 @@ async def _check_model_max_budget(
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and prisma_client is not None
and valid_token.token is not None
):
await model_max_budget_limiter.is_key_within_model_budget(

20
tests/test_litellm/proxy/auth/test_auth_model_budget_bypass.py
View file

@ -156,6 +156,7 @@ class TestCheckModelMaxBudget:
request_data={"model": "gpt-4"},
route="/chat/completions",
model_max_budget_limiter=limiter,
prisma_client=MagicMock(),
)
limiter.is_key_within_model_budget.assert_awaited_once_with(
user_api_key_dict=token,
@ -238,6 +239,7 @@ class TestCheckModelMaxBudget:
request_data={"model": "gpt-4"},
route="/chat/completions",
model_max_budget_limiter=limiter,
prisma_client=MagicMock(),
)
@pytest.mark.asyncio
@ -258,10 +260,28 @@ class TestCheckModelMaxBudget:
request_data={"model": "gpt-4"},
route="/chat/completions",
model_max_budget_limiter=limiter,
prisma_client=MagicMock(),
)
limiter.is_key_within_model_budget.assert_awaited_once()
limiter.is_end_user_within_model_budget.assert_awaited_once()
@pytest.mark.asyncio
async def test_should_skip_key_budget_when_prisma_client_is_none(self):
"""Key-level check requires prisma_client to be non-None."""
limiter = MagicMock(spec=_PROXY_VirtualKeyModelMaxBudgetLimiter)
token = UserAPIKeyAuth(
token="hashed-key",
model_max_budget={"gpt-4": {"budget_limit": 10.0, "time_period": "1d"}},
)
await _check_model_max_budget(
valid_token=token,
request_data={"model": "gpt-4"},
route="/chat/completions",
model_max_budget_limiter=limiter,
prisma_client=None,
)
limiter.is_key_within_model_budget.assert_not_called()
@pytest.mark.asyncio
async def test_should_skip_when_model_max_budget_empty(self):
"""Empty model_max_budget dict should be treated as no budget."""