From 5f9ea8a94f2811e8fcedfc1d3112d4b99901b160 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 2 Feb 2024 08:18:06 -0800
Subject: [PATCH] fix(utils.py): add track cost callback to callback list for
 team id callbacks

---
 litellm/proxy/proxy_server.py | 20 +++++++++-----------
 litellm/proxy/utils.py        |  3 ---
 litellm/utils.py              | 10 ++++++----
 3 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index c97ebc2b09..a55ba5e876 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -679,11 +679,11 @@ def cost_tracking():
     if prisma_client is not None or custom_db_client is not None:
         if isinstance(litellm.success_callback, list):
             verbose_proxy_logger.debug("setting litellm success callback to track cost")
-            if (track_cost_callback) not in litellm.success_callback:  # type: ignore
-                litellm.success_callback.append(track_cost_callback)  # type: ignore
+            if (_PROXY_track_cost_callback) not in litellm.success_callback:  # type: ignore
+                litellm.success_callback.append(_PROXY_track_cost_callback)  # type: ignore
 
 
-async def track_cost_callback(
+async def _PROXY_track_cost_callback(
     kwargs,  # kwargs to completion
     completion_response: litellm.ModelResponse,  # response from completion
     start_time=None,
@@ -752,8 +752,8 @@ async def update_database(
     end_time=None,
 ):
     try:
-        verbose_proxy_logger.debug(
-            f"Enters prisma db call, token: {token}; user_id: {user_id}"
+        verbose_proxy_logger.info(
+            f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}"
         )
 
         ### UPDATE USER SPEND ###
@@ -865,18 +865,16 @@ async def update_database(
             )
 
             payload["spend"] = response_cost
-
             if prisma_client is not None:
                 await prisma_client.insert_data(data=payload, table_name="spend")
 
             elif custom_db_client is not None:
                 await custom_db_client.insert_data(payload, table_name="spend")
 
-        tasks = []
-        tasks.append(_update_user_db())
-        tasks.append(_update_key_db())
-        tasks.append(_insert_spend_log_to_db())
-        await asyncio.gather(*tasks)
+        asyncio.create_task(_update_user_db())
+        asyncio.create_task(_update_key_db())
+        asyncio.create_task(_insert_spend_log_to_db())
+        verbose_proxy_logger.info("Successfully updated spend in all 3 tables")
     except Exception as e:
         verbose_proxy_logger.debug(
             f"Error updating Prisma database: {traceback.format_exc()}"
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 1a6e40319e..905b9424e1 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -524,7 +524,6 @@ class PrismaClient:
                     response = await self.db.litellm_verificationtoken.find_many(
                         order={"spend": "desc"},
                     )
-                print_verbose(f"PrismaClient: response={response}")
                 if response is not None:
                     return response
                 else:
@@ -1202,8 +1201,6 @@ async def reset_budget(prisma_client: PrismaClient):
             table_name="user", query_type="find_all", reset_at=now
         )
 
-        verbose_proxy_logger.debug(f"users_to_reset from get_data: {users_to_reset}")
-
         if users_to_reset is not None and len(users_to_reset) > 0:
             for user in users_to_reset:
                 user.spend = 0.0
diff --git a/litellm/utils.py b/litellm/utils.py
index 96b6bc0ac1..c087f98a24 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1470,10 +1470,12 @@ class Logging:
             callbacks = self.dynamic_async_success_callbacks
             ## keep the internal functions ##
             for callback in litellm._async_success_callback:
-                if (
-                    isinstance(callback, CustomLogger)
-                    and "_PROXY_" in callback.__class__.__name__
-                ):
+                callback_name = ""
+                if isinstance(callback, CustomLogger):
+                    callback_name = callback.__class__.__name__
+                if callable(callback):
+                    callback_name = callback.__name__
+                if "_PROXY_" in callback_name:
                     callbacks.append(callback)
         else:
             callbacks = litellm._async_success_callback