scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

Renamed param budget_reset_at to initial_budget_reset_at

Browse source
This commit is contained in:
Aayush Malviya 2026-03-02 14:20:00 +05:30
parent 001717bc07
commit 09a13621c6
4 changed files with 33 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/_types.py
View file

@ -786,7 +786,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
rpm_limit: Optional[int] = None
budget_duration: Optional[str] = None
budget_reset_at: Optional[datetime] = None
initial_budget_reset_at: Optional[datetime] = None
allowed_cache_controls: Optional[list] = []
config: Optional[dict] = {}
permissions: Optional[dict] = {}
@ -804,18 +804,18 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
@model_validator(mode="after")
def validate_budget_reset_at(self) -> "GenerateRequestBase":
"""Validate that budget_reset_at is not in the past"""
if self.budget_reset_at is not None:
def validate_initial_budget_reset_at(self) -> "GenerateRequestBase":
"""Validate that initial_budget_reset_at is not in the past"""
if self.initial_budget_reset_at is not None:
# Ensure timezone-aware for comparison
reset_at = self.budget_reset_at
reset_at = self.initial_budget_reset_at
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"budget_reset_at cannot be in the past. "
f"initial_budget_reset_at cannot be in the past. "
f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}"
)
return self
@ -874,6 +874,7 @@ class GenerateKeyResponse(KeyRequestBase):
key: str # type: ignore
key_name: Optional[str] = None
expires: Optional[datetime] = None
budget_reset_at: Optional[datetime] = None # computed/stored value from DB
user_id: Optional[str] = None
token_id: Optional[str] = None
organization_id: Optional[str] = None

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

@ -98,6 +98,10 @@ 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:
data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at")
return data_json
@ -306,7 +310,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").
- 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.
- 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)
@ -423,6 +427,9 @@ async def new_user(
response_dict["key"] = response.get("token", "")
if data.initial_budget_reset_at is not None:
response_dict["initial_budget_reset_at"] = data.initial_budget_reset_at
new_user_response = NewUserResponse(**response_dict)
#########################################################
@ -764,8 +771,8 @@ def _update_internal_user_params(
if data.user_role == LitellmUserRoles.INTERNAL_USER:
is_internal_user = True
if "budget_reset_at" in non_default_values and non_default_values["budget_reset_at"] is not None:
_reset_at = non_default_values["budget_reset_at"]
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):
@ -776,6 +783,8 @@ def _update_internal_user_params(
except Exception:
_reset_at = None
non_default_values["budget_reset_at"] = _reset_at
elif "initial_budget_reset_at" in non_default_values:
non_default_values.pop("initial_budget_reset_at")
if "budget_duration" in non_default_values:
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
@ -1003,6 +1012,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)

3
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -582,6 +582,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
if "budget_duration" in data_json:
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
if "initial_budget_reset_at" in data_json:
data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at", None)
if user_api_key_dict.user_id is not None:
data_json["created_by"] = user_api_key_dict.user_id
data_json["updated_by"] = user_api_key_dict.user_id

20
tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py → tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py
View file

@ -1,5 +1,5 @@
"""
Test cases for budget_reset_at feature in user management endpoints.
Test cases for initial_budget_reset_at feature in user management endpoints.
"""
import os
@ -19,20 +19,20 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
class TestBudgetResetAtValidation:
"""Test Pydantic validation for budget_reset_at field"""
"""Test Pydantic validation for initial_budget_reset_at field"""
def test_past_date_rejection_user_request(self):
"""Test that NewUserRequest rejects past dates for budget_reset_at"""
"""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",
budget_reset_at=past_date,
initial_budget_reset_at=past_date,
)
assert "budget_reset_at cannot be in the past" in str(exc_info.value)
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"""
@ -41,18 +41,18 @@ class TestBudgetResetAtValidation:
user_request = NewUserRequest(
user_id="test_user",
budget_duration="10m",
budget_reset_at=future_date,
initial_budget_reset_at=future_date,
)
assert user_request.budget_reset_at == future_date
assert user_request.initial_budget_reset_at == future_date
class TestUserCreationWithBudgetResetAt:
"""Integration test for creating users with budget_reset_at"""
"""Integration test for creating users with initial_budget_reset_at"""
@pytest.mark.asyncio
async def test_create_user_explicit_budget_reset_at_takes_precedence(self):
"""Test that explicit budget_reset_at is honored and takes precedence over duration-based computation"""
"""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=[])
@ -64,7 +64,7 @@ class TestUserCreationWithBudgetResetAt:
user_request = NewUserRequest(
user_id="new_user",
budget_duration="10m",
budget_reset_at=explicit_date,
initial_budget_reset_at=explicit_date,
max_budget=100.0,
)