From 3810b050c1956eea724be90b333dc2b21c3f0aae Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 29 Mar 2024 20:02:31 -0700
Subject: [PATCH] fix(proxy_server.py): increment cached global proxy spend
 object

---
 litellm/proxy/auth/auth_checks.py         |  2 +-
 litellm/proxy/proxy_server.py             | 17 +++++++++++++++--
 litellm/tests/test_key_generate_prisma.py | 15 +++++++++------
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 794d393c10..e3b7df4e9c 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -84,7 +84,7 @@ def common_checks(
     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}"
+                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 e5d57c8a64..0d2bd2bc23 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -472,6 +472,7 @@ async def user_api_key_auth(
                                 user_info=user_info,
                             )
                         )
+
                 # run through common checks
                 _ = common_checks(
                     request_body=request_data,
@@ -949,7 +950,6 @@ async def user_api_key_auth(
                             user_info=user_info,
                         )
                     )
-
             _ = common_checks(
                 request_body=request_data,
                 team_object=_team_obj,
@@ -1623,7 +1623,7 @@ async def update_cache(
 
     async def _update_user_cache():
         ## UPDATE CACHE FOR USER ID + GLOBAL PROXY
-        user_ids = [user_id, litellm_proxy_budget_name]
+        user_ids = [user_id]
         try:
             for _id in user_ids:
                 # Fetch the existing cost for the given user
@@ -1664,6 +1664,19 @@ async def update_cache(
                     user_api_key_cache.set_cache(
                         key=_id, value=existing_spend_obj.json()
                     )
+            ## UPDATE GLOBAL PROXY ##
+            global_proxy_spend = await user_api_key_cache.async_get_cache(
+                key="{}:spend".format(litellm_proxy_admin_name)
+            )
+            if global_proxy_spend is None:
+                await user_api_key_cache.async_set_cache(
+                    key="{}:spend".format(litellm_proxy_admin_name), value=response_cost
+                )
+            elif response_cost is not None and global_proxy_spend is not None:
+                increment = global_proxy_spend + response_cost
+                await user_api_key_cache.async_set_cache(
+                    key="{}:spend".format(litellm_proxy_admin_name), value=increment
+                )
         except Exception as e:
             verbose_proxy_logger.debug(
                 f"An error occurred updating user cache: {str(e)}\n\n{traceback.format_exc()}"
diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py
index 29468c449b..98126f9a16 100644
--- a/litellm/tests/test_key_generate_prisma.py
+++ b/litellm/tests/test_key_generate_prisma.py
@@ -409,18 +409,20 @@ def test_call_with_proxy_over_budget(prisma_client):
     litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}"
     setattr(
         litellm.proxy.proxy_server,
-        "litellm_proxy_budget_name",
+        "litellm_proxy_admin_name",
         litellm_proxy_budget_name,
     )
+    setattr(litellm, "max_budget", 0.00001)
+    from litellm.proxy.proxy_server import user_api_key_cache
+
+    user_api_key_cache.set_cache(
+        key="{}:spend".format(litellm_proxy_budget_name), value=0
+    )
+    setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
     try:
 
         async def test():
             await litellm.proxy.proxy_server.prisma_client.connect()
-            ## CREATE PROXY + USER BUDGET ##
-            request = NewUserRequest(
-                max_budget=0.00001, user_id=litellm_proxy_budget_name
-            )
-            await new_user(request)
             request = NewUserRequest()
             key = await new_user(request)
             print(key)
@@ -472,6 +474,7 @@ def test_call_with_proxy_over_budget(prisma_client):
                 start_time=datetime.now(),
                 end_time=datetime.now(),
             )
+
             await asyncio.sleep(5)
             # use generated key to auth in
             result = await user_api_key_auth(request=request, api_key=bearer_token)