From 16f3d7e0edd9e243bfd2c716cc58970731bcf670 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 18 Jan 2024 17:51:48 -0800
Subject: [PATCH] (feat) use user_id passed to request - cost track

---
 litellm/proxy/proxy_server.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 5d55a7161b..9e6589c482 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -523,19 +523,27 @@ async def track_cost_callback(
         verbose_proxy_logger.debug(
             f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}"
         )
+        litellm_params = kwargs.get("litellm_params", {})
+        proxy_server_request = litellm_params.get("proxy_server_request")
+        user_id = proxy_server_request.get("body", {}).get("user", None)
         if "complete_streaming_response" in kwargs:
             # for tracking streaming cost we pass the "messages" and the output_text to litellm.completion_cost
             completion_response = kwargs["complete_streaming_response"]
             response_cost = litellm.completion_cost(
                 completion_response=completion_response
             )
-            verbose_proxy_logger.debug(f"streaming response_cost {response_cost}")
+
             user_api_key = kwargs["litellm_params"]["metadata"].get(
                 "user_api_key", None
             )
-            user_id = kwargs["litellm_params"]["metadata"].get(
+
+            user_id = user_id or kwargs["litellm_params"]["metadata"].get(
                 "user_api_key_user_id", None
             )
+
+            verbose_proxy_logger.debug(
+                f"streaming response_cost {response_cost}, for user_id {user_id}"
+            )
             if user_api_key and (
                 prisma_client is not None or custom_db_client is not None
             ):
@@ -555,9 +563,12 @@ async def track_cost_callback(
             user_api_key = kwargs["litellm_params"]["metadata"].get(
                 "user_api_key", None
             )
-            user_id = kwargs["litellm_params"]["metadata"].get(
+            user_id = user_id or kwargs["litellm_params"]["metadata"].get(
                 "user_api_key_user_id", None
             )
+            verbose_proxy_logger.debug(
+                f"response_cost {response_cost}, for user_id {user_id}"
+            )
             if user_api_key and (
                 prisma_client is not None or custom_db_client is not None
             ):