mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge ff421f6a2c into 62757ff48f
This commit is contained in:
aayush-malviya-acquia
GitHub
commit
663e63ed2a
4 changed files with 277 additions and 11 deletions
|
|
@ -1,6 +1,6 @@
|
|||
import enum
|
||||
import json
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -1406,10 +1406,25 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
|
|||
user_role: Literal["user", "admin"] = "user"
|
||||
|
||||
|
||||
def _assert_initial_budget_reset_at_valid(value: Optional[datetime]) -> None:
|
||||
"""Assert that initial_budget_reset_at is not in the past. Raises ValueError if invalid."""
|
||||
if value is not None:
|
||||
reset_at = value
|
||||
if reset_at.tzinfo is None:
|
||||
reset_at = reset_at.replace(tzinfo=timezone.utc)
|
||||
current_time = datetime.now(timezone.utc)
|
||||
if reset_at < current_time:
|
||||
raise ValueError(
|
||||
f"initial_budget_reset_at cannot be in the past. "
|
||||
f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}"
|
||||
)
|
||||
|
||||
|
||||
class NewUserRequest(GenerateRequestBase):
|
||||
max_budget: Optional[float] = None
|
||||
user_email: Optional[str] = None
|
||||
user_alias: Optional[str] = None
|
||||
initial_budget_reset_at: Optional[datetime] = None
|
||||
user_role: Optional[
|
||||
Literal[
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
|
|
@ -1426,6 +1441,11 @@ class NewUserRequest(GenerateRequestBase):
|
|||
sso_user_id: Optional[str] = None
|
||||
organizations: Optional[List[str]] = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_initial_budget_reset_at(self) -> "NewUserRequest":
|
||||
_assert_initial_budget_reset_at_valid(self.initial_budget_reset_at)
|
||||
return self
|
||||
|
||||
|
||||
class NewUserResponse(GenerateKeyResponse):
|
||||
max_budget: Optional[float] = None
|
||||
|
|
@ -1443,6 +1463,7 @@ class NewUserResponse(GenerateKeyResponse):
|
|||
model_max_budget: Optional[dict] = None
|
||||
created_at: Optional[datetime] = None
|
||||
updated_at: Optional[datetime] = None
|
||||
initial_budget_reset_at: Optional[datetime] = None
|
||||
|
||||
|
||||
class UpdateUserRequestNoUserIDorEmail(
|
||||
|
|
@ -1461,6 +1482,12 @@ class UpdateUserRequestNoUserIDorEmail(
|
|||
]
|
||||
] = None
|
||||
max_budget: Optional[float] = None
|
||||
initial_budget_reset_at: Optional[datetime] = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_initial_budget_reset_at(self) -> "UpdateUserRequestNoUserIDorEmail":
|
||||
_assert_initial_budget_reset_at_valid(self.initial_budget_reset_at)
|
||||
return self
|
||||
|
||||
|
||||
class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail):
|
||||
|
|
|
|||
|
|
@ -119,6 +119,12 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d
|
|||
data_json["budget_duration"] = litellm.internal_user_budget_duration
|
||||
|
||||
data_json.pop("teams", None) # handled separately
|
||||
|
||||
if "initial_budget_reset_at" in data_json and data_json.get("initial_budget_reset_at") is not None:
|
||||
data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at")
|
||||
else:
|
||||
data_json.pop("initial_budget_reset_at", None)
|
||||
|
||||
return data_json
|
||||
|
||||
|
||||
|
|
@ -369,6 +375,7 @@ async def new_user(
|
|||
- user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`
|
||||
- max_budget: Optional[float] - Specify max budget for a given user.
|
||||
- budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").
|
||||
- initial_budget_reset_at: Optional[datetime] - Specify the exact datetime when the budget should first reset (ISO 8601 format, e.g., "2025-11-15T00:00:00Z"). If provided along with budget_duration, this value takes precedence for the initial reset; subsequent resets will follow the budget_duration interval.
|
||||
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models). Set to ['no-default-models'] to block all model access. Restricting user to only team-based model access.
|
||||
- tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute)
|
||||
- rpm_limit: Optional[int] - Specify rpm limit for a given user (Requests per minute)
|
||||
|
|
@ -503,6 +510,9 @@ async def new_user(
|
|||
|
||||
response_dict["key"] = response.get("token", "")
|
||||
|
||||
if data.initial_budget_reset_at is not None and response.get("budget_reset_at") is not None:
|
||||
response_dict["initial_budget_reset_at"] = response.get("budget_reset_at")
|
||||
|
||||
new_user_response = NewUserResponse(**response_dict)
|
||||
|
||||
#########################################################
|
||||
|
|
@ -1086,12 +1096,29 @@ def _update_internal_user_params(
|
|||
if data.user_role == LitellmUserRoles.INTERNAL_USER:
|
||||
is_internal_user = True
|
||||
|
||||
if "initial_budget_reset_at" in non_default_values and non_default_values["initial_budget_reset_at"] is not None:
|
||||
_reset_at = non_default_values.pop("initial_budget_reset_at")
|
||||
try:
|
||||
# convert to datetime if not already
|
||||
if not isinstance(_reset_at, datetime):
|
||||
_reset_at = datetime.fromisoformat(str(_reset_at).replace("Z", "+00:00"))
|
||||
# ensure timezone aware
|
||||
if _reset_at.tzinfo is None:
|
||||
_reset_at = _reset_at.replace(tzinfo=timezone.utc)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to parse initial_budget_reset_at value: {_reset_at}. Error: {e}. Falling back to None."
|
||||
)
|
||||
_reset_at = None
|
||||
non_default_values["budget_reset_at"] = _reset_at
|
||||
|
||||
if "budget_duration" in non_default_values:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
non_default_values["budget_reset_at"] = get_budget_reset_time(
|
||||
budget_duration=non_default_values["budget_duration"]
|
||||
)
|
||||
# only auto-compute budget_reset_at if not explicitly provided
|
||||
if "budget_reset_at" not in non_default_values or non_default_values["budget_reset_at"] is None:
|
||||
non_default_values["budget_reset_at"] = get_budget_reset_time(
|
||||
budget_duration=non_default_values["budget_duration"]
|
||||
)
|
||||
|
||||
if "max_budget" not in non_default_values:
|
||||
if (
|
||||
|
|
@ -1108,9 +1135,11 @@ def _update_internal_user_params(
|
|||
] = litellm.internal_user_budget_duration
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
non_default_values["budget_reset_at"] = get_budget_reset_time(
|
||||
budget_duration=non_default_values["budget_duration"]
|
||||
)
|
||||
# only auto-compute budget_reset_at if not explicitly provided
|
||||
if "budget_reset_at" not in non_default_values or non_default_values["budget_reset_at"] is None:
|
||||
non_default_values["budget_reset_at"] = get_budget_reset_time(
|
||||
budget_duration=non_default_values["budget_duration"]
|
||||
)
|
||||
|
||||
return non_default_values
|
||||
|
||||
|
|
@ -1312,6 +1341,7 @@ async def user_update(
|
|||
- user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`
|
||||
- max_budget: Optional[float] - Specify max budget for a given user.
|
||||
- budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo").
|
||||
- initial_budget_reset_at: Optional[datetime] - Specify the exact datetime when the budget should first reset (ISO 8601 format, e.g., "2025-11-15T00:00:00Z"). If provided along with budget_duration, this value takes precedence for the initial reset; subsequent resets will follow the budget_duration interval.
|
||||
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
|
||||
- tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute)
|
||||
- rpm_limit: Optional[int] - Specify rpm limit for a given user (Requests per minute)
|
||||
|
|
|
|||
|
|
@ -2801,6 +2801,7 @@ async def generate_key_helper_fn( # noqa: PLR0915
|
|||
max_budget: Optional[float] = None, # max_budget is used to Budget Per user
|
||||
blocked: Optional[bool] = None,
|
||||
budget_duration: Optional[str] = None, # max_budget is used to Budget Per user
|
||||
budget_reset_at: Optional[datetime] = None, # manual override for initial user/key budget reset time
|
||||
token: Optional[str] = None,
|
||||
key: Optional[
|
||||
str
|
||||
|
|
@ -2869,10 +2870,30 @@ async def generate_key_helper_fn( # noqa: PLR0915
|
|||
else:
|
||||
key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration)
|
||||
|
||||
if budget_duration is None: # one-time budget
|
||||
reset_at = None
|
||||
if budget_reset_at is not None:
|
||||
try:
|
||||
# try to parse budget_reset_at
|
||||
if isinstance(budget_reset_at, datetime):
|
||||
_reset_candidate = budget_reset_at
|
||||
else:
|
||||
_reset_candidate = datetime.fromisoformat(str(budget_reset_at).replace("Z", "+00:00"))
|
||||
# ensure tz-aware datetime
|
||||
if _reset_candidate.tzinfo is None:
|
||||
_reset_candidate = _reset_candidate.replace(tzinfo=timezone.utc)
|
||||
reset_at = _reset_candidate
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to parse budget_reset_at value: {budget_reset_at}. Error: {e}. Falling back to budget_duration-based computation."
|
||||
)
|
||||
if budget_duration is None: # fallback to one-time budget
|
||||
reset_at = None
|
||||
else:
|
||||
reset_at = get_budget_reset_time(budget_duration=budget_duration) # calculate based on budget duration
|
||||
else:
|
||||
reset_at = get_budget_reset_time(budget_duration=budget_duration)
|
||||
if budget_duration is None: # one-time budget
|
||||
reset_at = None
|
||||
else:
|
||||
reset_at = get_budget_reset_time(budget_duration=budget_duration) # calculate based on budget duration
|
||||
|
||||
aliases_json = json.dumps(aliases)
|
||||
config_json = json.dumps(config)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,188 @@
|
|||
"""
|
||||
Test cases for initial_budget_reset_at feature in user management endpoints.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../..")
|
||||
)
|
||||
|
||||
from litellm.proxy._types import NewUserRequest, UpdateUserRequest, UpdateUserRequestNoUserIDorEmail
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user, user_update
|
||||
|
||||
|
||||
class TestInitialBudgetResetAtValidation:
|
||||
"""Test Pydantic validation for initial_budget_reset_at field"""
|
||||
|
||||
def test_past_date_rejection_user_request(self):
|
||||
"""Test that NewUserRequest rejects past dates for initial_budget_reset_at"""
|
||||
past_date = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
NewUserRequest(
|
||||
user_id="test_user",
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=past_date,
|
||||
)
|
||||
|
||||
assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value)
|
||||
|
||||
def test_future_date_accepted(self):
|
||||
"""Test that future dates are accepted"""
|
||||
future_date = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
|
||||
user_request = NewUserRequest(
|
||||
user_id="test_user",
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=future_date,
|
||||
)
|
||||
|
||||
assert user_request.initial_budget_reset_at == future_date
|
||||
|
||||
def test_past_date_rejection_update_user_request(self):
|
||||
"""Test that UpdateUserRequest rejects past dates for initial_budget_reset_at"""
|
||||
past_date = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
UpdateUserRequest(
|
||||
user_id="test_user",
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=past_date,
|
||||
)
|
||||
|
||||
assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value)
|
||||
|
||||
def test_past_date_rejection_update_user_request_no_id_or_email(self):
|
||||
"""Test that UpdateUserRequestNoUserIDorEmail rejects past dates for initial_budget_reset_at"""
|
||||
past_date = datetime.now(timezone.utc) - timedelta(days=1)
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
UpdateUserRequestNoUserIDorEmail(
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=past_date,
|
||||
)
|
||||
|
||||
assert "initial_budget_reset_at cannot be in the past" in str(exc_info.value)
|
||||
|
||||
def test_future_date_accepted_update_user_request(self):
|
||||
"""Test that future dates are accepted on UpdateUserRequest"""
|
||||
future_date = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
|
||||
update_request = UpdateUserRequest(
|
||||
user_id="test_user",
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=future_date,
|
||||
)
|
||||
|
||||
assert update_request.initial_budget_reset_at == future_date
|
||||
|
||||
|
||||
class TestUserCreationWithInitialBudgetResetAt:
|
||||
"""Integration test for creating users with initial_budget_reset_at"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_user_explicit_initial_budget_reset_at_takes_precedence(self):
|
||||
"""Test that explicit initial_budget_reset_at is honored and takes precedence over duration-based computation"""
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.insert_data = AsyncMock(
|
||||
return_value=MagicMock(user_id="new_user", spend=0.0, models=[])
|
||||
)
|
||||
mock_prisma_client.db.litellm_usertable.count = AsyncMock(return_value=0)
|
||||
|
||||
explicit_date = datetime.now(timezone.utc) + timedelta(days=30)
|
||||
|
||||
user_request = NewUserRequest(
|
||||
user_id="new_user",
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=explicit_date,
|
||||
max_budget=100.0,
|
||||
)
|
||||
|
||||
mock_license_check = MagicMock()
|
||||
mock_license_check.is_over_limit.return_value = False
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
mock_prisma_client,
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server._license_check",
|
||||
mock_license_check,
|
||||
):
|
||||
result = await new_user(
|
||||
data=user_request,
|
||||
user_api_key_dict=MagicMock(user_id="admin", user_role="proxy_admin"),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert mock_prisma_client.insert_data.called
|
||||
|
||||
# find the user insert call and verify budget_reset_at
|
||||
for call in mock_prisma_client.insert_data.call_args_list:
|
||||
if call.kwargs.get("table_name") == "user":
|
||||
user_data = call.kwargs["data"]
|
||||
assert "budget_reset_at" in user_data
|
||||
stored_date = user_data["budget_reset_at"]
|
||||
assert isinstance(stored_date, datetime)
|
||||
assert abs((stored_date - explicit_date).total_seconds()) < 1 # verify dates are effectively equal
|
||||
break
|
||||
else:
|
||||
raise AssertionError("User insert call not found")
|
||||
|
||||
|
||||
class TestUserUpdateWithInitialBudgetResetAt:
|
||||
"""Integration test for updating users with initial_budget_reset_at"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_user_explicit_initial_budget_reset_at_translated(self):
|
||||
"""Test that initial_budget_reset_at is translated to budget_reset_at on the update path"""
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
|
||||
mock_prisma_client.update_data = AsyncMock(
|
||||
return_value={"user_id": "existing_user"}
|
||||
)
|
||||
|
||||
explicit_date = datetime.now(timezone.utc) + timedelta(days=30)
|
||||
|
||||
update_request = UpdateUserRequest(
|
||||
user_id="existing_user",
|
||||
budget_duration="10m",
|
||||
initial_budget_reset_at=explicit_date,
|
||||
max_budget=100.0,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
mock_prisma_client,
|
||||
), patch(
|
||||
"litellm.proxy.proxy_server.litellm_proxy_admin_name",
|
||||
"admin",
|
||||
):
|
||||
result = await user_update(
|
||||
data=update_request,
|
||||
user_api_key_dict=MagicMock(
|
||||
user_id="existing_user",
|
||||
user_role="proxy_admin",
|
||||
),
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert mock_prisma_client.update_data.called
|
||||
|
||||
call_kwargs = mock_prisma_client.update_data.call_args.kwargs
|
||||
user_data = call_kwargs["data"]
|
||||
assert "budget_reset_at" in user_data, "budget_reset_at should be set from initial_budget_reset_at"
|
||||
assert "initial_budget_reset_at" not in user_data, "initial_budget_reset_at should be removed"
|
||||
stored_date = user_data["budget_reset_at"]
|
||||
assert isinstance(stored_date, datetime)
|
||||
assert abs((stored_date - explicit_date).total_seconds()) < 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Loading…
Add table
Reference in a new issue