From a299ac2328d4ef32c25cae8e45bc3248a253bacd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Jan 2024 20:51:13 -0800 Subject: [PATCH 1/4] fix(utils.py): enable cost tracking for image gen models on proxy --- litellm/main.py | 2 +- litellm/tests/test_custom_callback_input.py | 68 ++++++++++----------- litellm/utils.py | 15 ++--- 3 files changed, 43 insertions(+), 42 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index f9f1139f69b..c809f49d694 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3067,7 +3067,7 @@ def image_generation( custom_llm_provider=custom_llm_provider, **non_default_params, ) - logging = litellm_logging_obj + logging: Logging = litellm_logging_obj logging.update_environment_variables( model=model, user=user, diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index a61cc843ec1..641343e7a2a 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -819,44 +819,44 @@ async def test_async_embedding_azure_caching(): # Image Generation -# ## Test OpenAI + Sync -# def test_image_generation_openai(): -# try: -# customHandler_success = CompletionCustomHandler() -# customHandler_failure = CompletionCustomHandler() -# litellm.callbacks = [customHandler_success] +## Test OpenAI + Sync +def test_image_generation_openai(): + try: + customHandler_success = CompletionCustomHandler() + customHandler_failure = CompletionCustomHandler() + litellm.callbacks = [customHandler_success] -# litellm.set_verbose = True + litellm.set_verbose = True -# response = litellm.image_generation( -# prompt="A cute baby sea otter", model="dall-e-3" -# ) + response = litellm.image_generation( + prompt="A cute baby sea otter", model="dall-e-3" + ) -# print(f"response: {response}") -# assert len(response.data) > 0 + print(f"response: {response}") + assert len(response.data) > 0 -# print(f"customHandler_success.errors: {customHandler_success.errors}") -# print(f"customHandler_success.states: {customHandler_success.states}") -# assert len(customHandler_success.errors) == 0 -# assert len(customHandler_success.states) == 3 # pre, post, success -# # test failure callback -# litellm.callbacks = [customHandler_failure] -# try: -# response = litellm.image_generation( -# prompt="A cute baby sea otter", model="dall-e-4" -# ) -# except: -# pass -# print(f"customHandler_failure.errors: {customHandler_failure.errors}") -# print(f"customHandler_failure.states: {customHandler_failure.states}") -# assert len(customHandler_failure.errors) == 0 -# assert len(customHandler_failure.states) == 3 # pre, post, failure -# except litellm.RateLimitError as e: -# pass -# except litellm.ContentPolicyViolationError: -# pass # OpenAI randomly raises these errors - skip when they occur -# except Exception as e: -# pytest.fail(f"An exception occurred - {str(e)}") + print(f"customHandler_success.errors: {customHandler_success.errors}") + print(f"customHandler_success.states: {customHandler_success.states}") + assert len(customHandler_success.errors) == 0 + assert len(customHandler_success.states) == 3 # pre, post, success + # test failure callback + litellm.callbacks = [customHandler_failure] + try: + response = litellm.image_generation( + prompt="A cute baby sea otter", model="dall-e-4" + ) + except: + pass + print(f"customHandler_failure.errors: {customHandler_failure.errors}") + print(f"customHandler_failure.states: {customHandler_failure.states}") + assert len(customHandler_failure.errors) == 0 + assert len(customHandler_failure.states) == 3 # pre, post, failure + except litellm.RateLimitError as e: + pass + except litellm.ContentPolicyViolationError: + pass # OpenAI randomly raises these errors - skip when they occur + except Exception as e: + pytest.fail(f"An exception occurred - {str(e)}") # test_image_generation_openai() diff --git a/litellm/utils.py b/litellm/utils.py index b0e48bbc6e2..613d9d90ae2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2029,14 +2029,15 @@ def client(original_function): start_time=start_time, ) ## check if metadata is passed in + litellm_params = {} if "metadata" in kwargs: - litellm_params = {"metadata": kwargs["metadata"]} - logging_obj.update_environment_variables( - model=model, - user="", - optional_params={}, - litellm_params=litellm_params, - ) + litellm_params["metadata"] = kwargs["metadata"] + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=litellm_params, + ) return logging_obj except Exception as e: import logging From dad578f96ace828951b334a4f2d3d25364e83374 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Jan 2024 20:52:38 -0800 Subject: [PATCH 2/4] build(schema.prisma): update schema --- litellm/proxy/schema.prisma | 7 ++++++- schema.prisma | 2 ++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2eb6332092a..02e4114e5d9 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -7,6 +7,7 @@ generator client { provider = "prisma-client-py" } +// Track spend, rate limit, budget Users model LiteLLM_UserTable { user_id String @unique team_id String? @@ -21,9 +22,11 @@ model LiteLLM_UserTable { budget_reset_at DateTime? } -// required for token gen +// Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @unique + key_name String? + key_alias String? spend Float @default(0.0) expires DateTime? models String[] @@ -40,11 +43,13 @@ model LiteLLM_VerificationToken { budget_reset_at DateTime? } +// store proxy config.yaml model LiteLLM_Config { param_name String @id param_value Json? } +// View spend, model, api_key per request model LiteLLM_SpendLogs { request_id String @unique call_type String diff --git a/schema.prisma b/schema.prisma index 0882c650c81..02e4114e5d9 100644 --- a/schema.prisma +++ b/schema.prisma @@ -25,6 +25,8 @@ model LiteLLM_UserTable { // Generate Tokens for Proxy model LiteLLM_VerificationToken { token String @unique + key_name String? + key_alias String? spend Float @default(0.0) expires DateTime? models String[] From 618d216c1123ec6f5bc8a951bb96ec308aba3f70 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Jan 2024 21:04:06 -0800 Subject: [PATCH 3/4] fix(openai.py): fix image gen logging --- litellm/tests/test_custom_callback_input.py | 28 +++++++++++---------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/litellm/tests/test_custom_callback_input.py b/litellm/tests/test_custom_callback_input.py index 641343e7a2a..266303df15a 100644 --- a/litellm/tests/test_custom_callback_input.py +++ b/litellm/tests/test_custom_callback_input.py @@ -824,26 +824,28 @@ def test_image_generation_openai(): try: customHandler_success = CompletionCustomHandler() customHandler_failure = CompletionCustomHandler() - litellm.callbacks = [customHandler_success] + # litellm.callbacks = [customHandler_success] - litellm.set_verbose = True + # litellm.set_verbose = True - response = litellm.image_generation( - prompt="A cute baby sea otter", model="dall-e-3" - ) + # response = litellm.image_generation( + # prompt="A cute baby sea otter", model="dall-e-3" + # ) - print(f"response: {response}") - assert len(response.data) > 0 + # print(f"response: {response}") + # assert len(response.data) > 0 - print(f"customHandler_success.errors: {customHandler_success.errors}") - print(f"customHandler_success.states: {customHandler_success.states}") - assert len(customHandler_success.errors) == 0 - assert len(customHandler_success.states) == 3 # pre, post, success + # print(f"customHandler_success.errors: {customHandler_success.errors}") + # print(f"customHandler_success.states: {customHandler_success.states}") + # assert len(customHandler_success.errors) == 0 + # assert len(customHandler_success.states) == 3 # pre, post, success # test failure callback litellm.callbacks = [customHandler_failure] try: response = litellm.image_generation( - prompt="A cute baby sea otter", model="dall-e-4" + prompt="A cute baby sea otter", + model="dall-e-2", + api_key="my-bad-api-key", ) except: pass @@ -859,7 +861,7 @@ def test_image_generation_openai(): pytest.fail(f"An exception occurred - {str(e)}") -# test_image_generation_openai() +test_image_generation_openai() ## Test OpenAI + Async ## Test Azure + Sync From d755d509013d2e1eef434f2a0365fdbe8e0c1dc8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Jan 2024 21:05:49 -0800 Subject: [PATCH 4/4] fix(openai.py): fix openai image gen logging --- litellm/llms/openai.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 01887616caf..da89b7796ef 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -718,8 +718,22 @@ class OpenAIChatCompletion(BaseLLM): return convert_to_model_response_object(response_object=response, model_response_object=model_response, response_type="image_generation") # type: ignore except OpenAIError as e: exception_mapping_worked = True + ## LOGGING + logging_obj.post_call( + input=prompt, + api_key=api_key, + additional_args={"complete_input_dict": data}, + original_response=str(e), + ) raise e except Exception as e: + ## LOGGING + logging_obj.post_call( + input=prompt, + api_key=api_key, + additional_args={"complete_input_dict": data}, + original_response=str(e), + ) if hasattr(e, "status_code"): raise OpenAIError(status_code=e.status_code, message=str(e)) else: