From b5c5ed220910370b73937208d2e2f481eb966ec9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 24 Jul 2024 15:02:03 -0700 Subject: [PATCH] fix(key_management_endpoints.py): if budget duration set, set budget_reset_at --- .../proxy/management_endpoints/key_management_endpoints.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 03028821f..0e4696e44 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -333,6 +333,13 @@ async def update_key_fn( expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s) non_default_values["expires"] = expires + if "budget_duration" in non_default_values: + duration_s = _duration_in_seconds( + duration=non_default_values["budget_duration"] + ) + key_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + non_default_values["budget_reset_at"] = key_reset_at + response = await prisma_client.update_data( token=key, data={**non_default_values, "token": key} )