From 9609df16d308a203811ed0fb460d03b8d7c19f14 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 22 May 2024 13:16:37 -0700
Subject: [PATCH] feat(proxy_server.py): new 'add budget' endpoint

create a budget object (max budget, tpm /rpm limits, etc.) and assign that to a user/team/etc.
---
 litellm/proxy/_types.py       | 25 +++++++++++++++++++++++
 litellm/proxy/proxy_server.py | 37 +++++++++++++++++++++++++++++++++--
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index a6640e4d3..8bfa56004 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -676,6 +676,31 @@ class OrganizationRequest(LiteLLMBase):
     organizations: List[str]
 
 
+class BudgetNew(LiteLLMBase):
+    budget_id: str = Field(default=None, description="The unique budget id.")
+    max_budget: Optional[float] = Field(
+        default=None,
+        description="Requests will fail if this budget (in USD) is exceeded.",
+    )
+    soft_budget: Optional[float] = Field(
+        default=None,
+        description="Requests will NOT fail if this is exceeded. Will fire alerting though.",
+    )
+    max_parallel_requests: Optional[int] = Field(
+        default=None, description="Max concurrent requests allowed for this budget id."
+    )
+    tpm_limit: Optional[int] = Field(
+        default=None, description="Max tokens per minute, allowed for this budget id."
+    )
+    rpm_limit: Optional[int] = Field(
+        default=None, description="Max requests per minute, allowed for this budget id."
+    )
+    budget_duration: Optional[str] = Field(
+        default=None,
+        description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')",
+    )
+
+
 class BudgetRequest(LiteLLMBase):
     budgets: List[str]
 
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index c7bf50926..f157f420c 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -3470,7 +3470,9 @@ async def startup_event():
     if prisma_client is not None and master_key is not None:
         # add master key to db
         if os.getenv("PROXY_ADMIN_ID", None) is not None:
-            litellm_proxy_admin_name = os.getenv("PROXY_ADMIN_ID")
+            litellm_proxy_admin_name = os.getenv(
+                "PROXY_ADMIN_ID", litellm_proxy_admin_name
+            )
         asyncio.create_task(
             generate_key_helper_fn(
                 duration=None,
@@ -8187,9 +8189,40 @@ async def info_organization(data: OrganizationRequest):
 #### BUDGET TABLE MANAGEMENT ####
 
 
+@router.post(
+    "/budget/new",
+    tags=["budget management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def new_budget(
+    budget_obj: BudgetNew,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Create a new budget object. Can apply this to teams, orgs, end-users, keys.
+    """
+    global prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(
+            status_code=500,
+            detail={"error": CommonProxyErrors.db_not_connected_error.value},
+        )
+
+    response = await prisma_client.db.litellm_budgettable.create(
+        data={
+            **budget_obj.model_dump(exclude_none=True),  # type: ignore
+            "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+            "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+        }  # type: ignore
+    )
+
+    return response
+
+
 @router.post(
     "/budget/info",
-    tags=["organization management"],
+    tags=["budget management"],
     dependencies=[Depends(user_api_key_auth)],
 )
 async def info_budget(data: BudgetRequest):