From 1872122bfb4694919c4e3ba9b9609a5f1ad660c4 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 11:26:40 -0800 Subject: [PATCH] (test) streaming spend logs test - actually run this --- litellm/tests/test_key_generate_prisma.py | 123 ++++++++++------------ 1 file changed, 56 insertions(+), 67 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index fe5f8fa5654..f130807c680 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -774,7 +774,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") @@ -784,81 +785,69 @@ 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={ - "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) - - 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) - 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))