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

fix(utils.py): return last streaming chunk

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-29 12:11:08 -08:00
parent a9ed768991
commit 2b437a2699
2 changed files with 8 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

14
litellm/tests/test_custom_logger.py
View file

@ -37,15 +37,7 @@ def test_chat_openai():
"role": "user", "role": "user",
"content": "Hi 👋 - i'm openai" "content": "Hi 👋 - i'm openai"
}], }],
stream=True, stream=True)
complete_response = True)
response2 = completion(model="gpt-3.5-turbo",
messages=[{
"role": "user",
"content": "Hi 👋 - i'm not openai"
}],
stream=True,
complete_response = True)
time.sleep(1) time.sleep(1)
assert customHandler.success == True assert customHandler.success == True
except Exception as e: except Exception as e:
@ -53,7 +45,7 @@ def test_chat_openai():
pass pass
# test_chat_openai() test_chat_openai()
def test_completion_azure_stream_moderation_failure(): def test_completion_azure_stream_moderation_failure():
try: try:
@ -80,7 +72,7 @@ def test_completion_azure_stream_moderation_failure():
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_completion_azure_stream_moderation_failure() # test_completion_azure_stream_moderation_failure()
# def custom_callback( # def custom_callback(

10
litellm/utils.py
View file

@ -5242,6 +5242,11 @@ class CustomStreamWrapper:
return model_response return model_response
else: else:
return return
elif model_response.choices[0].finish_reason:
model_response.choices[0].finish_reason = map_finish_reason(model_response.choices[0].finish_reason) # ensure consistent output to openai
# LOGGING
threading.Thread(target=self.logging_obj.success_handler, args=(model_response,)).start()
return model_response
elif response_obj is not None and response_obj.get("original_chunk", None) is not None: # function / tool calling branch - only set for openai/azure compatible endpoints elif response_obj is not None and response_obj.get("original_chunk", None) is not None: # function / tool calling branch - only set for openai/azure compatible endpoints
# enter this branch when no content has been passed in response # enter this branch when no content has been passed in response
original_chunk = response_obj.get("original_chunk", None) original_chunk = response_obj.get("original_chunk", None)
@ -5263,11 +5268,6 @@ class CustomStreamWrapper:
self.sent_first_chunk = True self.sent_first_chunk = True
threading.Thread(target=self.logging_obj.success_handler, args=(model_response,)).start() # log response threading.Thread(target=self.logging_obj.success_handler, args=(model_response,)).start() # log response
return model_response return model_response
elif model_response.choices[0].finish_reason:
model_response.choices[0].finish_reason = map_finish_reason(model_response.choices[0].finish_reason) # ensure consistent output to openai
# LOGGING
threading.Thread(target=self.logging_obj.success_handler, args=(model_response,)).start()
return model_response
else: else:
return return
except StopIteration: except StopIteration: