From 2776d9cb0da7e4fad8b675bfe733c69817417b9b Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 8 Aug 2024 13:05:03 -0700
Subject: [PATCH] fix(internal_user_endpoints.py): expose new
 'internal_user_budget_duration' flag

Relevant to - https://github.com/BerriAI/litellm/issues/5106
---
 litellm/__init__.py                                      | 1 +
 litellm/proxy/_new_secret_config.yaml                    | 9 +++++++++
 .../management_endpoints/internal_user_endpoints.py      | 4 ++++
 3 files changed, 14 insertions(+)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index 0429c654748..8d0e1b31495 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -261,6 +261,7 @@ default_user_params: Optional[Dict] = None
 default_team_settings: Optional[List] = None
 max_user_budget: Optional[float] = None
 max_internal_user_budget: Optional[float] = None
+internal_user_budget_duration: Optional[str] = None
 max_end_user_budget: Optional[float] = None
 #### REQUEST PRIORITIZATION ####
 priority_reservation: Optional[Dict[str, float]] = None
diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
index 97457ef84b2..0388b4b94a6 100644
--- a/litellm/proxy/_new_secret_config.yaml
+++ b/litellm/proxy/_new_secret_config.yaml
@@ -2,3 +2,12 @@ model_list:
   - model_name: "gpt-4"
     litellm_params:
       model: "gpt-4"
+      api_key: "bad_key"
+  - model_name: "gpt-4o"
+    litellm_params:
+      model: "gpt-4o"
+
+litellm_settings:
+  max_internal_user_budget: 0.001
+  internal_user_budget_duration: "5m"
+
diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index f54eca9a2fc..173b205fc30 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -91,6 +91,10 @@ async def new_user(
         if litellm.max_internal_user_budget is not None:
             data_json["max_budget"] = litellm.max_internal_user_budget
 
+    if "budget_duration" in data_json and data_json["budget_duration"] is None:
+        if litellm.internal_user_budget_duration is not None:
+            data_json["budget_duration"] = litellm.internal_user_budget_duration
+
     response = await generate_key_helper_fn(request_type="user", **data_json)
 
     # Admin UI Logic