phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(utils.py): fix streaming on-success logging

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-28 09:11:47 -08:00
parent 6750fa6d9a
commit 6a7a8be4fb
2 changed files with 99 additions and 100 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/main.py
View file

@ -2176,16 +2176,15 @@ def stream_chunk_builder(chunks: list, messages: Optional[list]=None):
content_list.append(content) content_list.append(content)
# Combine the "content" strings into a single string || combine the 'function' strings into a single string # 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 # Update the "content" field within the response dictionary
response["choices"][0]["message"]["content"] = combined_content response["choices"][0]["message"]["content"] = combined_content
if len(combined_content) > 0: if len(combined_content) > 0:
completion_output = combined_content completion_output = combined_content
elif len(combined_arguments) > 0: elif len(combined_arguments) > 0:
completion_output = combined_arguments completion_output = combined_arguments
# # Update usage information if needed # # Update usage information if needed
if messages: if messages:
response["usage"]["prompt_tokens"] = token_counter(model=model, messages=messages) response["usage"]["prompt_tokens"] = token_counter(model=model, messages=messages)

194
litellm/tests/test_custom_logger.py
View file

@ -1,122 +1,122 @@
# ### What this tests #### ### What this tests ####
# import sys import sys, os, time
# import os 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 class MyCustomHandler(CustomLogger):
# import litellm success: bool = False
# from litellm.integrations.custom_logger import CustomLogger
# class MyCustomHandler(CustomLogger): def log_pre_api_call(self, model, messages, kwargs):
# def log_pre_api_call(self, model, messages, kwargs): print(f"Pre-API Call")
# print(f"Pre-API Call")
# def log_post_api_call(self, kwargs, response_obj, start_time, end_time): def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
# print(f"Post-API Call") print(f"Post-API Call")
# def log_stream_event(self, kwargs, response_obj, start_time, end_time): def log_stream_event(self, kwargs, response_obj, start_time, end_time):
# print(f"On Stream") print(f"On Stream")
# def log_success_event(self, kwargs, response_obj, start_time, end_time): def log_success_event(self, kwargs, response_obj, start_time, end_time):
# print(f"On Success") print(f"On Success")
self.success = True
# def log_failure_event(self, kwargs, response_obj, start_time, end_time): def log_failure_event(self, kwargs, response_obj, start_time, end_time):
# print(f"On Failure") print(f"On Failure")
# customHandler = MyCustomHandler() def test_chat_openai():
try:
# def test_chat_openai(): customHandler = MyCustomHandler()
# try: litellm.callbacks = [customHandler]
# litellm.callbacks = [customHandler] response = completion(model="gpt-3.5-turbo",
# response = completion(model="gpt-3.5-turbo", messages=[{
# messages=[{ "role": "user",
# "role": "user", "content": "Hi 👋 - i'm openai"
# "content": "Hi 👋 - i'm openai" }],
# }], stream=True,
# stream=True) complete_response = True)
# for chunk in response:
# # print(chunk) time.sleep(1)
# continue assert customHandler.success == True
# # print(response) except Exception as e:
pytest.fail(f"An error occurred - {str(e)}")
# except Exception as e: pass
# print(e)
# pass
# test_chat_openai() test_chat_openai()
# # def custom_callback( # def custom_callback(
# # kwargs, # kwargs,
# # completion_response, # completion_response,
# # start_time, # start_time,
# # end_time, # end_time,
# # ): # ):
# # print( # print(
# # "in custom callback func" # "in custom callback func"
# # ) # )
# # print("kwargs", kwargs) # print("kwargs", kwargs)
# # print(completion_response) # print(completion_response)
# # print(start_time) # print(start_time)
# # print(end_time) # print(end_time)
# # if "complete_streaming_response" in kwargs: # if "complete_streaming_response" in kwargs:
# # print("\n\n complete response\n\n") # print("\n\n complete response\n\n")
# # complete_streaming_response = kwargs["complete_streaming_response"] # complete_streaming_response = kwargs["complete_streaming_response"]
# # print(kwargs["complete_streaming_response"]) # print(kwargs["complete_streaming_response"])
# # usage = complete_streaming_response["usage"] # usage = complete_streaming_response["usage"]
# # print("usage", usage) # print("usage", usage)
# # def send_slack_alert( # def send_slack_alert(
# # kwargs, # kwargs,
# # completion_response, # completion_response,
# # start_time, # start_time,
# # end_time, # end_time,
# # ): # ):
# # print( # print(
# # "in custom slack callback func" # "in custom slack callback func"
# # ) # )
# # import requests # import requests
# # import json # import json
# # # Define the Slack webhook URL # # Define the Slack webhook URL
# # slack_webhook_url = os.environ['SLACK_WEBHOOK_URL'] # "https://hooks.slack.com/services/<>/<>/<>" # slack_webhook_url = os.environ['SLACK_WEBHOOK_URL'] # "https://hooks.slack.com/services/<>/<>/<>"
# # # Define the text payload, send data available in litellm custom_callbacks # # 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)} # 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 = { # payload = {
# # "text": text_payload # "text": text_payload
# # } # }
# # # Set the headers # # Set the headers
# # headers = { # headers = {
# # "Content-type": "application/json" # "Content-type": "application/json"
# # } # }
# # # Make the POST request # # Make the POST request
# # response = requests.post(slack_webhook_url, json=payload, headers=headers) # response = requests.post(slack_webhook_url, json=payload, headers=headers)
# # # Check the response status # # Check the response status
# # if response.status_code == 200: # if response.status_code == 200:
# # print("Message sent successfully to Slack!") # print("Message sent successfully to Slack!")
# # else: # else:
# # print(f"Failed to send message to Slack. Status code: {response.status_code}") # print(f"Failed to send message to Slack. Status code: {response.status_code}")
# # print(response.json()) # print(response.json())
# # def get_transformed_inputs( # def get_transformed_inputs(
# # kwargs, # kwargs,
# # ): # ):
# # params_to_model = kwargs["additional_args"]["complete_input_dict"] # params_to_model = kwargs["additional_args"]["complete_input_dict"]
# # print("params to model", params_to_model) # print("params to model", params_to_model)
# # litellm.success_callback = [custom_callback, send_slack_alert] # litellm.success_callback = [custom_callback, send_slack_alert]
# # litellm.failure_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]