From 786116783f118e7316239bd2f18e6fd2d92d4926 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 29 Mar 2024 16:24:40 -0700
Subject: [PATCH] fix(proxy_server.py): fix max budget check to also fire slack
 alert

---
 litellm/proxy/_new_secret_config.yaml |  4 +++
 litellm/proxy/auth/auth_checks.py     | 13 ++++-----
 litellm/proxy/proxy_server.py         | 42 ++++++++++++++++++++++-----
 3 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
index 6159536c29e..d218fddb19d 100644
--- a/litellm/proxy/_new_secret_config.yaml
+++ b/litellm/proxy/_new_secret_config.yaml
@@ -5,6 +5,10 @@ model_list:
     api_key: my-fake-key
     api_base: https://exampleopenaiendpoint-production.up.railway.app/
 
+litellm_settings:
+  max_budget: 600020
+  budget_duration: 30d
+
 general_settings:
   master_key: sk-1234
   proxy_batch_write_at: 60 # 👈 Frequency of batch writing logs to server (in seconds)
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 6f95fad15eb..e9280fd31d3 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -81,14 +81,11 @@ def common_checks(
                 f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
             )
     # 6. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
-    if (
-        litellm.max_budget > 0
-        and global_proxy_spend is not None
-        and global_proxy_spend > litellm.max_budget
-    ):
-        raise Exception(
-            f"'ExceededBudget: LiteLLM Proxy has exceeded its budget. Current spend: {global_proxy_spend}; Max Budget: {litellm.max_budget}"
-        )
+    if litellm.max_budget > 0 and global_proxy_spend is not None:
+        if global_proxy_spend > litellm.max_budget:
+            raise Exception(
+                f"'ExceededBudget: LiteLLM Proxy has exceeded its budget. Current spend: {global_proxy_spend}; Max Budget: {litellm.max_budget}"
+            )
     return True
 
 
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 2c8106f2ab4..15615797604 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -446,7 +446,7 @@ async def user_api_key_auth(
                     )
                     if global_proxy_spend is None and prisma_client is not None:
                         # get from db
-                        sql_query = """SELECT SUM(spend) as total_spend FROM MONTHLYGLOBALSPEND;"""
+                        sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
 
                         response = await prisma_client.db.query_raw(query=sql_query)
 
@@ -457,7 +457,21 @@ async def user_api_key_auth(
                             value=global_proxy_spend,
                             ttl=60,
                         )
-
+                    if global_proxy_spend is not None:
+                        user_info = {
+                            "user_id": litellm_proxy_admin_name,
+                            "max_budget": litellm.max_budget,
+                            "spend": global_proxy_spend,
+                            "user_email": "",
+                        }
+                        asyncio.create_task(
+                            proxy_logging_obj.budget_alerts(
+                                user_max_budget=litellm.max_budget,
+                                user_current_spend=global_proxy_spend,
+                                type="user_and_proxy_budget",
+                                user_info=user_info,
+                            )
+                        )
                 # run through common checks
                 _ = common_checks(
                     request_body=request_data,
@@ -903,20 +917,33 @@ async def user_api_key_auth(
                 )
                 if global_proxy_spend is None:
                     # get from db
-                    sql_query = (
-                        """SELECT SUM(spend) as total_spend FROM MONTHLYGLOBALSPEND;"""
-                    )
+                    sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
 
                     response = await prisma_client.db.query_raw(query=sql_query)
 
-                    global_proxy_spend = response[0].total_spend
-
+                    global_proxy_spend = response[0]["total_spend"]
                     await user_api_key_cache.async_set_cache(
                         key="{}:spend".format(litellm_proxy_admin_name),
                         value=global_proxy_spend,
                         ttl=60,
                     )
 
+                if global_proxy_spend is not None:
+                    user_info = {
+                        "user_id": litellm_proxy_admin_name,
+                        "max_budget": litellm.max_budget,
+                        "spend": global_proxy_spend,
+                        "user_email": "",
+                    }
+                    asyncio.create_task(
+                        proxy_logging_obj.budget_alerts(
+                            user_max_budget=litellm.max_budget,
+                            user_current_spend=global_proxy_spend,
+                            type="user_and_proxy_budget",
+                            user_info=user_info,
+                        )
+                    )
+
             _ = common_checks(
                 request_body=request_data,
                 team_object=_team_obj,
@@ -4086,7 +4113,6 @@ async def generate_key_fn(
             )
             _budget_id = getattr(_budget, "budget_id", None)
         data_json = data.json()  # type: ignore
-
         # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
         if "max_budget" in data_json:
             data_json["key_max_budget"] = data_json.pop("max_budget", None)