Merge pull request #1591 from BerriAI/litellm_test_spend_logging

[Test] Proxy - writing spend logs table

litellm/proxy/utils.py

@@ -892,6 +892,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time):
     from pydantic import Json
     import uuid
 
+    verbose_proxy_logger.debug(
+        f"SpendTable: get_logging_payload - kwargs: {kwargs}\n\n"
+    )
+
     if kwargs == None:
         kwargs = {}
     # standardize this function to be used across, s3, dynamoDB, langfuse logging

litellm/tests/test_key_generate_prisma.py

@@ -42,6 +42,7 @@ from litellm.proxy.proxy_server import (
     info_key_fn,
     update_key_fn,
     generate_key_fn,
+    view_spend_logs,
 )
 from litellm.proxy.utils import PrismaClient, ProxyLogging
 from litellm._logging import verbose_proxy_logger
@@ -713,9 +714,12 @@ def test_call_with_key_over_budget(prisma_client):
             # update spend using track_cost callback, make 2nd request, it should fail
             from litellm.proxy.proxy_server import track_cost_callback
             from litellm import ModelResponse, Choices, Message, Usage
+            import time
+
+            request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"
 
             resp = ModelResponse(
-                id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
+                id=request_id,
                 choices=[
                     Choices(
                         finish_reason=None,
@@ -731,6 +735,7 @@ def test_call_with_key_over_budget(prisma_client):
             )
             await track_cost_callback(
                 kwargs={
+                    "model": "chatgpt-v-2",
                     "stream": False,
                     "litellm_params": {
                         "metadata": {
@@ -745,6 +750,18 @@ def test_call_with_key_over_budget(prisma_client):
                 end_time=datetime.now(),
             )
 
+            # test spend_log was written and we can read it
+            spend_logs = await view_spend_logs(request_id=request_id)
+
+            print("read spend logs", spend_logs)
+            assert len(spend_logs) == 1
+
+            spend_log = spend_logs[0]
+
+            assert spend_log.request_id == request_id
+            assert spend_log.spend == float("2e-05")
+            assert spend_log.model == "chatgpt-v-2"
+
             # use generated key to auth in
             result = await user_api_key_auth(request=request, api_key=bearer_token)
             print("result from user auth with new key", result)
@@ -788,9 +805,11 @@ def test_call_with_key_over_budget_stream(prisma_client):
             # update spend using track_cost callback, make 2nd request, it should fail
             from litellm.proxy.proxy_server import track_cost_callback
             from litellm import ModelResponse, Choices, Message, Usage
+            import time
 
+            request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"
             resp = ModelResponse(
-                id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
+                id=request_id,
                 choices=[
                     Choices(
                         finish_reason=None,
@@ -806,6 +825,7 @@ def test_call_with_key_over_budget_stream(prisma_client):
             )
             await track_cost_callback(
                 kwargs={
+                    "model": "sagemaker-chatgpt-v-2",
                     "stream": True,
                     "complete_streaming_response": resp,
                     "litellm_params": {
@@ -814,13 +834,25 @@ def test_call_with_key_over_budget_stream(prisma_client):
                             "user_api_key_user_id": user_id,
                         }
                     },
-                    "response_cost": 0.00002,
+                    "response_cost": 0.00005,
                 },
                 completion_response=ModelResponse(),
                 start_time=datetime.now(),
                 end_time=datetime.now(),
             )
 
+            # test spend_log was written and we can read it
+            spend_logs = await view_spend_logs(request_id=request_id)
+
+            print("read spend logs", spend_logs)
+            assert len(spend_logs) == 1
+
+            spend_log = spend_logs[0]
+
+            assert spend_log.request_id == request_id
+            assert spend_log.spend == float("5e-05")
+            assert spend_log.model == "sagemaker-chatgpt-v-2"
+
             # use generated key to auth in
             result = await user_api_key_auth(request=request, api_key=bearer_token)
             print("result from user auth with new key", result)