add cost tracking

litellm/integrations/custom_logger.py

@@ -3,6 +3,8 @@
 import dotenv, os
 import requests
 import requests
+import inspect
+import asyncio
 
 dotenv.load_dotenv()  # Loading env variables using dotenv
 import traceback
@@ -50,16 +52,27 @@ class CustomLogger:
         # Method definition
         try:
             kwargs["log_event_type"] = "post_api_call"
-            callback_func(
-                kwargs, # kwargs to func
-                response_obj,
-                start_time,
-                end_time,
-            )
+            if inspect.iscoroutinefunction(callback_func):
+                # If it's async, use asyncio to run it
+
+                loop = asyncio.get_event_loop()
+                if loop.is_closed():
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                loop.run_until_complete(callback_func(kwargs, response_obj, start_time, end_time))
+            else:
+                # If it's not async, run it synchronously
+                callback_func(
+                    kwargs,  # kwargs to func
+                    response_obj,
+                    start_time,
+                    end_time,
+                )
             print_verbose(
                 f"Custom Logger - final response object: {response_obj}"
             )
-        except:
+        except Exception as e:
+            raise e
             # traceback.print_exc()
             print_verbose(f"Custom Logger Error - {traceback.format_exc()}")
             pass