fix(proxy_server.py): increment cached global proxy spend object

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
 

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()}"

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)