scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

fix(proxy): preserve model budget update compatibility

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-11 22:59:46 -07:00
parent f96af80a0c
commit f44d5ef101
No known key found for this signature in database
2 changed files with 29 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -22,7 +22,7 @@ from typing import Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -88,6 +88,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import (
SCIM_ENTITLEMENTS_METADATA_KEY,
SCIM_ROLES_METADATA_KEY,
)
from litellm.types.utils import BudgetConfig
if TYPE_CHECKING:
from prisma import models as prisma_models
@ -98,7 +99,8 @@ if TYPE_CHECKING:
from litellm.proxy.utils import ProxyLogging
router: Final = APIRouter()
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(GenericBudgetConfigType)
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
def _user_table(
@ -1257,7 +1259,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
non_default_values[k] = v
elif k == "model_max_budget":
if k in fields_set:
_USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v)
try:
_USER_MODEL_BUDGET_ADAPTER.validate_python({} if v is None else v)
except ValidationError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
non_default_values[k] = {} if v is None else v
elif (
v is not None
@ -1888,10 +1893,13 @@ async def bulk_user_update(
)
if "model_max_budget" in non_default_values:
await evict_and_broadcast(
cache_keys=tuple(user.user_id for user in all_users_in_db),
user_api_key_cache=user_api_key_cache,
)
for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE):
await asyncio.gather(
*(
evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache)
for user in all_users_in_db[start : start + _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE]
)
)
# Create individual success results
for user in all_users_in_db:

17
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -5,7 +5,7 @@ from typing import Final
import pytest
from fastapi.testclient import TestClient
from pydantic import ValidationError
from fastapi import HTTPException
from pytest_mock import MockerFixture
@ -2147,7 +2147,7 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo
"user_alias": "Spruce",
}
replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}}
replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0}
request: Final = UpdateUserRequest(
user_id="user-spruce",
model_max_budget=replacement,
@ -2168,8 +2168,9 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo
def test_update_internal_user_params_rejects_invalid_model_budget(invalid_budget: dict[str, object]) -> None:
request: Final = UpdateUserRequest(user_id="user-spruce", model_max_budget=invalid_budget)
with pytest.raises(ValidationError):
with pytest.raises(HTTPException) as exc:
_update_internal_user_params(data_json=request.model_dump(exclude_unset=True), data=request)
assert exc.value.status_code == 400
@pytest.mark.asyncio
@ -2228,6 +2229,16 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
new_callable=mocker.AsyncMock,
)
with pytest.raises(HTTPException) as exc:
await bulk_user_update(
data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": {"model-spruce": "invalid"}}),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),
litellm_changed_by=None,
)
assert exc.value.status_code == 400
prisma_client.db.litellm_usertable.update_many.assert_not_called()
assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) == saved_user
response: Final = await bulk_user_update(
data=BulkUpdateUserRequest(all_users=True, user_updates={"model_max_budget": None}),
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),