fix(key_management_endpoints.py): if budget duration set, set budget_…

…reset_at
BerriAI · Jul 24, 2024 · b5c5ed2 · b5c5ed2
1 parent dc3b39c
commit b5c5ed2
Showing 1 changed file with 7 additions and 0 deletions.
diff --git 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}
         )