fix(internal_user_endpoints.py): support updating budgets for /user/update

This commit is contained in:
Krrish Dholakia 2024-07-24 21:51:24 -07:00
parent 780f9a78f1
commit d1622f6c0c

View file

@ -27,6 +27,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import * from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import (
_duration_in_seconds,
generate_key_helper_fn, generate_key_helper_fn,
) )
from litellm.proxy.management_helpers.utils import ( from litellm.proxy.management_helpers.utils import (
@ -486,6 +487,13 @@ async def user_update(
): # models default to [], spend defaults to 0, we should not reset these values ): # models default to [], spend defaults to 0, we should not reset these values
non_default_values[k] = v non_default_values[k] = v
if "budget_duration" in non_default_values:
duration_s = _duration_in_seconds(
duration=non_default_values["budget_duration"]
)
user_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
non_default_values["budget_reset_at"] = user_reset_at
## ADD USER, IF NEW ## ## ADD USER, IF NEW ##
verbose_proxy_logger.debug("/user/update: Received data = %s", data) verbose_proxy_logger.debug("/user/update: Received data = %s", data)
if data.user_id is not None and len(data.user_id) > 0: if data.user_id is not None and len(data.user_id) > 0: