phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(test) track usage in custom callback streaming

Browse source
This commit is contained in:
ishaan-jaff 2023-10-31 23:05:42 -07:00
parent d1f2593dc0
commit f6983223f9
1 changed files with 13 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/tests/test_custom_logger.py
View file

@ -20,6 +20,12 @@ def custom_callback(
print(completion_response) print(completion_response)
print(start_time) print(start_time)
print(end_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( def send_slack_alert(
kwargs, kwargs,
completion_response, completion_response,
@ -67,20 +73,23 @@ litellm.success_callback = [custom_callback, send_slack_alert]
litellm.failure_callback = [send_slack_alert] litellm.failure_callback = [send_slack_alert]
litellm.set_verbose = True litellm.set_verbose = False
litellm.input_callback = [get_transformed_inputs] # litellm.input_callback = [get_transformed_inputs]
def test_chat_openai(): def test_chat_openai():
try: try:
response = completion(model="gpt-2", 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)
print(response) print(response)
for chunk in response:
print(chunk)
except Exception as e: except Exception as e:
print(e) print(e)