From 09a13621c6b31ac573e45a687f327552b0c3658e Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Mon, 2 Mar 2026 14:20:00 +0530 Subject: [PATCH] Renamed param budget_reset_at to initial_budget_reset_at --- litellm/proxy/_types.py | 13 ++++++------ .../internal_user_endpoints.py | 16 ++++++++++++--- .../key_management_endpoints.py | 3 +++ ..._at.py => test_initial_budget_reset_at.py} | 20 +++++++++---------- 4 files changed, 33 insertions(+), 19 deletions(-) rename tests/test_litellm/proxy/management_endpoints/{test_budget_reset_at.py => test_initial_budget_reset_at.py} (78%) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 30f24dcde7a..d05a2ce0b43 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 681087451bc..1026b967a8f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 74ed92ce1f8..d0715eae959 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py b/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py similarity index 78% rename from tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py rename to tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py index a63f00d8c99..40875b13736 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_reset_at.py +++ b/tests/test_litellm/proxy/management_endpoints/test_initial_budget_reset_at.py @@ -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, )