feat - reset spend per team, api_key

enterprise/utils.py

@@ -44,8 +44,8 @@ async def ui_get_spend_by_tags(start_date: str, end_date: str, prisma_client):
     # print("tags - spend")
     # print(response)
     # Bar Chart 1 - Spend per tag - Top 10 tags by spend
-    total_spend_per_tag = collections.defaultdict(float)
-    total_requests_per_tag = collections.defaultdict(int)
+    total_spend_per_tag: collections.defaultdict = collections.defaultdict(float)
+    total_requests_per_tag: collections.defaultdict = collections.defaultdict(int)
     for row in response:
         tag_name = row["individual_request_tag"]
         tag_spend = row["total_spend"]

litellm/proxy/proxy_server.py

@@ -5922,6 +5922,63 @@ async def view_spend_logs(
         )
 
 
+@router.get(
+    "/global/spend/reset",
+    tags=["Budget & Spend Tracking"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def global_spend_reset(
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    ADMIN ONLY / MASTER KEY Only Endpoint
+
+    Globally reset spend for All API Keys and Teams, maintain LiteLLM_SpendLogs
+
+    1. LiteLLM_SpendLogs will maintain the logs on spend, no data gets deleted from there
+    2. LiteLLM_VerificationTokens spend will be set = 0
+    3. LiteLLM_TeamTable spend will be set = 0
+
+    """
+    global prisma_client, master_key
+    if prisma_client is None:
+        raise ProxyException(
+            message="Prisma Client is not initialized",
+            type="internal_error",
+            param="None",
+            code=status.HTTP_401_UNAUTHORIZED,
+        )
+
+    if master_key is None:
+        raise ProxyException(
+            message="Master key is not initialized, please set LITELLM_MASTER_KEY in .env",
+            type="internal_error",
+            param="None",
+            code=status.HTTP_401_UNAUTHORIZED,
+        )
+
+    if user_api_key_dict.api_key is None:
+        raise ProxyException(
+            message="no api_key passed",
+            type="auth_error",
+            param="master_key",
+            code=status.HTTP_401_UNAUTHORIZED,
+        )
+
+    if not secrets.compare_digest(master_key, user_api_key_dict.api_key):
+        raise ProxyException(
+            message="/global/spend/reset Route only allowed for master key",
+            type="auth_error",
+            param="master_key",
+            code=status.HTTP_401_UNAUTHORIZED,
+        )
+
+    await prisma_client.db.litellm_verificationtoken.update_many(
+        data={"spend": 0.0}, where={}
+    )
+    await prisma_client.db.litellm_teamtable.update_many(data={"spend": 0.0}, where={})
+
+
 @router.get(
     "/global/spend/logs",
     tags=["Budget & Spend Tracking"],