From 150b91d476a447d97cd21d28b7131423bafd933d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 28 Nov 2023 09:11:47 -0800 Subject: [PATCH] fix(utils.py): fix streaming on-success logging --- litellm/main.py | 5 +- litellm/tests/test_custom_logger.py | 194 ++++++++++++++-------------- 2 files changed, 99 insertions(+), 100 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 28464b5208..3f096535d6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2176,16 +2176,15 @@ def stream_chunk_builder(chunks: list, messages: Optional[list]=None): content_list.append(content) # Combine the "content" strings into a single string || combine the 'function' strings into a single string - combined_content = "".join(combined_arguments) + combined_content = "".join(content_list) # Update the "content" field within the response dictionary response["choices"][0]["message"]["content"] = combined_content - + if len(combined_content) > 0: completion_output = combined_content elif len(combined_arguments) > 0: completion_output = combined_arguments - # # Update usage information if needed if messages: response["usage"]["prompt_tokens"] = token_counter(model=model, messages=messages) diff --git a/litellm/tests/test_custom_logger.py b/litellm/tests/test_custom_logger.py index ab0458e7cd..b6c0894c44 100644 --- a/litellm/tests/test_custom_logger.py +++ b/litellm/tests/test_custom_logger.py @@ -1,122 +1,122 @@ -# ### What this tests #### -# import sys -# import os +### What this tests #### +import sys, os, time +import pytest +sys.path.insert(0, os.path.abspath('../..')) -# sys.path.insert(0, os.path.abspath('../..')) +from litellm import completion, embedding +import litellm +from litellm.integrations.custom_logger import CustomLogger -# from litellm import completion, embedding -# import litellm -# from litellm.integrations.custom_logger import CustomLogger +class MyCustomHandler(CustomLogger): + success: bool = False -# class MyCustomHandler(CustomLogger): -# def log_pre_api_call(self, model, messages, kwargs): -# print(f"Pre-API Call") + def log_pre_api_call(self, model, messages, kwargs): + print(f"Pre-API Call") -# def log_post_api_call(self, kwargs, response_obj, start_time, end_time): -# print(f"Post-API Call") + def log_post_api_call(self, kwargs, response_obj, start_time, end_time): + print(f"Post-API Call") -# def log_stream_event(self, kwargs, response_obj, start_time, end_time): -# print(f"On Stream") + def log_stream_event(self, kwargs, response_obj, start_time, end_time): + print(f"On Stream") -# def log_success_event(self, kwargs, response_obj, start_time, end_time): -# print(f"On Success") + def log_success_event(self, kwargs, response_obj, start_time, end_time): + print(f"On Success") + self.success = True -# def log_failure_event(self, kwargs, response_obj, start_time, end_time): -# print(f"On Failure") + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + print(f"On Failure") -# customHandler = MyCustomHandler() - -# def test_chat_openai(): -# try: -# litellm.callbacks = [customHandler] -# response = completion(model="gpt-3.5-turbo", -# messages=[{ -# "role": "user", -# "content": "Hi 👋 - i'm openai" -# }], -# stream=True) -# for chunk in response: -# # print(chunk) -# continue -# # print(response) - -# except Exception as e: -# print(e) -# pass +def test_chat_openai(): + try: + customHandler = MyCustomHandler() + litellm.callbacks = [customHandler] + response = completion(model="gpt-3.5-turbo", + messages=[{ + "role": "user", + "content": "Hi 👋 - i'm openai" + }], + stream=True, + complete_response = True) + + time.sleep(1) + assert customHandler.success == True + except Exception as e: + pytest.fail(f"An error occurred - {str(e)}") + pass -# test_chat_openai() +test_chat_openai() -# # def custom_callback( -# # kwargs, -# # completion_response, -# # start_time, -# # end_time, -# # ): -# # print( -# # "in custom callback func" -# # ) -# # print("kwargs", kwargs) -# # print(completion_response) -# # print(start_time) -# # print(end_time) -# # if "complete_streaming_response" in kwargs: -# # print("\n\n complete response\n\n") -# # complete_streaming_response = kwargs["complete_streaming_response"] -# # print(kwargs["complete_streaming_response"]) -# # usage = complete_streaming_response["usage"] -# # print("usage", usage) -# # def send_slack_alert( -# # kwargs, -# # completion_response, -# # start_time, -# # end_time, -# # ): -# # print( -# # "in custom slack callback func" -# # ) -# # import requests -# # import json +# def custom_callback( +# kwargs, +# completion_response, +# start_time, +# end_time, +# ): +# print( +# "in custom callback func" +# ) +# print("kwargs", kwargs) +# print(completion_response) +# print(start_time) +# print(end_time) +# if "complete_streaming_response" in kwargs: +# print("\n\n complete response\n\n") +# complete_streaming_response = kwargs["complete_streaming_response"] +# print(kwargs["complete_streaming_response"]) +# usage = complete_streaming_response["usage"] +# print("usage", usage) +# def send_slack_alert( +# kwargs, +# completion_response, +# start_time, +# end_time, +# ): +# print( +# "in custom slack callback func" +# ) +# import requests +# import json -# # # Define the Slack webhook URL -# # slack_webhook_url = os.environ['SLACK_WEBHOOK_URL'] # "https://hooks.slack.com/services/<>/<>/<>" +# # Define the Slack webhook URL +# slack_webhook_url = os.environ['SLACK_WEBHOOK_URL'] # "https://hooks.slack.com/services/<>/<>/<>" -# # # Define the text payload, send data available in litellm custom_callbacks -# # text_payload = f"""LiteLLM Logging: kwargs: {str(kwargs)}\n\n, response: {str(completion_response)}\n\n, start time{str(start_time)} end time: {str(end_time)} -# # """ -# # payload = { -# # "text": text_payload -# # } +# # Define the text payload, send data available in litellm custom_callbacks +# text_payload = f"""LiteLLM Logging: kwargs: {str(kwargs)}\n\n, response: {str(completion_response)}\n\n, start time{str(start_time)} end time: {str(end_time)} +# """ +# payload = { +# "text": text_payload +# } -# # # Set the headers -# # headers = { -# # "Content-type": "application/json" -# # } +# # Set the headers +# headers = { +# "Content-type": "application/json" +# } -# # # Make the POST request -# # response = requests.post(slack_webhook_url, json=payload, headers=headers) +# # Make the POST request +# response = requests.post(slack_webhook_url, json=payload, headers=headers) -# # # Check the response status -# # if response.status_code == 200: -# # print("Message sent successfully to Slack!") -# # else: -# # print(f"Failed to send message to Slack. Status code: {response.status_code}") -# # print(response.json()) +# # Check the response status +# if response.status_code == 200: +# print("Message sent successfully to Slack!") +# else: +# print(f"Failed to send message to Slack. Status code: {response.status_code}") +# print(response.json()) -# # def get_transformed_inputs( -# # kwargs, -# # ): -# # params_to_model = kwargs["additional_args"]["complete_input_dict"] -# # print("params to model", params_to_model) +# def get_transformed_inputs( +# kwargs, +# ): +# params_to_model = kwargs["additional_args"]["complete_input_dict"] +# print("params to model", params_to_model) -# # litellm.success_callback = [custom_callback, send_slack_alert] -# # litellm.failure_callback = [send_slack_alert] +# litellm.success_callback = [custom_callback, send_slack_alert] +# litellm.failure_callback = [send_slack_alert] -# # litellm.set_verbose = False +# litellm.set_verbose = False -# # # litellm.input_callback = [get_transformed_inputs] +# # litellm.input_callback = [get_transformed_inputs]