From 9116c09386314f5a746e5b9649290b82b4de0330 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 4 Nov 2024 11:33:27 -0800 Subject: [PATCH] fix test key gen prisma --- .../local_testing/test_key_generate_prisma.py | 219 +++++++++++------- 1 file changed, 137 insertions(+), 82 deletions(-) diff --git a/tests/local_testing/test_key_generate_prisma.py b/tests/local_testing/test_key_generate_prisma.py index 74182c09f..0cb9659dc 100644 --- a/tests/local_testing/test_key_generate_prisma.py +++ b/tests/local_testing/test_key_generate_prisma.py @@ -105,6 +105,12 @@ from litellm.proxy._types import ( UpdateUserRequest, UserAPIKeyAuth, ) +from litellm.types.utils import ( + StandardLoggingPayload, + StandardLoggingModelInformation, + StandardLoggingMetadata, + StandardLoggingHiddenParams, +) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) @@ -141,6 +147,58 @@ def prisma_client(): return prisma_client +def create_simple_standard_logging_payload() -> StandardLoggingPayload: + + return StandardLoggingPayload( + id="test_id", + call_type="completion", + response_cost=0.1, + response_cost_failure_debug_info=None, + status="success", + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + startTime=1234567890.0, + endTime=1234567891.0, + completionStartTime=1234567890.5, + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ), + model="gpt-3.5-turbo", + model_id="model-123", + model_group="openai-gpt", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_org_id=None, + user_api_key_alias="test_alias", + user_api_key_team_id="test_team", + user_api_key_user_id="test_user", + user_api_key_team_alias="test_team_alias", + spend_logs_metadata=None, + requester_ip_address="127.0.0.1", + requester_metadata=None, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address="127.0.0.1", + messages=[{"role": "user", "content": "Hello, world!"}], + response={"choices": [{"message": {"content": "Hi there!"}}]}, + error_str=None, + model_parameters={"stream": True}, + hidden_params=StandardLoggingHiddenParams( + model_id="model-123", + cache_key=None, + api_base="https://api.openai.com", + response_cost="0.1", + additional_headers=None, + ), + ) + + @pytest.mark.asyncio() @pytest.mark.flaky(retries=6, delay=1) async def test_new_user_response(prisma_client): @@ -514,16 +572,17 @@ def test_call_with_user_over_budget(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id + await track_cost_callback( kwargs={ "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": generated_key, - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -611,21 +670,17 @@ def test_call_with_end_user_over_budget(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 10 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token="sk-1234" + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user + standard_logging_payload["end_user"] = user await track_cost_callback( kwargs={ "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": "sk-1234", - "user_api_key_user_id": user, - }, - "proxy_server_request": { - "body": { - "user": user, - } - }, - }, - "response_cost": 10, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -717,16 +772,16 @@ def test_call_with_proxy_over_budget(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": generated_key, - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -808,17 +863,17 @@ def test_call_with_user_over_budget_stream(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "stream": True, "complete_streaming_response": resp, - "litellm_params": { - "metadata": { - "user_api_key": generated_key, - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=ModelResponse(), start_time=datetime.now(), @@ -914,17 +969,17 @@ def test_call_with_proxy_over_budget_stream(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "stream": True, "complete_streaming_response": resp, - "litellm_params": { - "metadata": { - "user_api_key": generated_key, - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=ModelResponse(), start_time=datetime.now(), @@ -1471,17 +1526,17 @@ def test_call_with_key_over_budget(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "model": "chatgpt-v-2", "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -1588,17 +1643,17 @@ def test_call_with_key_over_budget_no_cache(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "model": "chatgpt-v-2", "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -1712,17 +1767,17 @@ def test_call_with_key_over_model_budget(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00002 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "model": "chatgpt-v-2", "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -1818,17 +1873,17 @@ async def test_call_with_key_never_over_budget(prisma_client): prompt_tokens=210000, completion_tokens=200000, total_tokens=41000 ), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 200000 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id await track_cost_callback( kwargs={ "model": "chatgpt-v-2", "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 200000, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -1899,19 +1954,19 @@ async def test_call_with_key_over_budget_stream(prisma_client): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00005 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id 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": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00005, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(), @@ -2292,19 +2347,19 @@ async def track_cost_callback_helper_fn(generated_key: str, user_id: str): model="gpt-35-turbo", # azure always has model written like this usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), ) + standard_logging_payload = create_simple_standard_logging_payload() + standard_logging_payload["response_cost"] = 0.00005 + standard_logging_payload["metadata"]["user_api_key_hash"] = hash_token( + token=generated_key + ) + standard_logging_payload["metadata"]["user_api_key_user_id"] = user_id 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": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00005, + "standard_logging_object": standard_logging_payload, }, completion_response=resp, start_time=datetime.now(),