(test) streaming spend logs test

litellm/tests/test_key_generate_prisma.py

@@ -775,7 +775,8 @@ def test_call_with_key_over_budget(prisma_client):
         print(vars(e))
 
 
-def test_call_with_key_over_budget_stream(prisma_client):
+@pytest.mark.asyncio()
+async def test_call_with_key_over_budget_stream(prisma_client):
     # 14. Make a call with a key over budget, expect to fail
     setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
     setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
@@ -785,82 +786,81 @@ def test_call_with_key_over_budget_stream(prisma_client):
     litellm.set_verbose = True
     verbose_proxy_logger.setLevel(logging.DEBUG)
     try:
+        await litellm.proxy.proxy_server.prisma_client.connect()
+        request = GenerateKeyRequest(max_budget=0.00001)
+        key = await generate_key_fn(request)
+        print(key)
 
-        async def test():
-            await litellm.proxy.proxy_server.prisma_client.connect()
-            request = GenerateKeyRequest(max_budget=0.00001)
-            key = await generate_key_fn(request)
-            print(key)
+        generated_key = key.key
+        user_id = key.user_id
+        bearer_token = "Bearer " + generated_key
 
-            generated_key = key.key
-            user_id = key.user_id
-            bearer_token = "Bearer " + generated_key
+        request = Request(scope={"type": "http"})
+        request._url = URL(url="/chat/completions")
 
-            request = Request(scope={"type": "http"})
-            request._url = URL(url="/chat/completions")
+        # 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)
 
-            # 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)
+        # 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
 
-            # 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=request_id,
-                choices=[
-                    Choices(
-                        finish_reason=None,
-                        index=0,
-                        message=Message(
-                            content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
-                            role="assistant",
-                        ),
-                    )
-                ],
-                model="gpt-35-turbo",  # azure always has model written like this
-                usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
-            )
-            await track_cost_callback(
-                kwargs={
-                    "call_type": "acompletion",
-                    "model": "sagemaker-chatgpt-v-2",
-                    "stream": True,
-                    "complete_streaming_response": resp,
-                    "litellm_params": {
-                        "metadata": {
-                            "user_api_key": generated_key,
-                            "user_api_key_user_id": user_id,
-                        }
-                    },
-                    "response_cost": 0.00005,
+        request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{time.time()}"
+        resp = ModelResponse(
+            id=request_id,
+            choices=[
+                Choices(
+                    finish_reason=None,
+                    index=0,
+                    message=Message(
+                        content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
+                        role="assistant",
+                    ),
+                )
+            ],
+            model="gpt-35-turbo",  # azure always has model written like this
+            usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410),
+        )
+        await track_cost_callback(
+            kwargs={
+                "call_type": "acompletion",
+                "model": "sagemaker-chatgpt-v-2",
+                "stream": True,
+                "complete_streaming_response": resp,
+                "litellm_params": {
+                    "metadata": {
+                        "user_api_key": generated_key,
+                        "user_api_key_user_id": user_id,
+                    }
                 },
-                completion_response=ModelResponse(),
-                start_time=datetime.now(),
-                end_time=datetime.now(),
-            )
+                "response_cost": 0.00005,
+            },
+            completion_response=resp,
+            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)
+        # 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
+        print("read spend logs", spend_logs)
+        assert len(spend_logs) == 1
 
-            spend_log = spend_logs[0]
+        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"
+        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)
-            pytest.fail(f"This should have failed!. They key crossed it's budget")
+        # 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)
+        pytest.fail(f"This should have failed!. They key crossed it's budget")
 
     except Exception as e:
+        print("Got Exception", e)
         error_detail = e.message
         assert "Authentication Error, ExceededTokenBudget:" in error_detail
         print(vars(e))