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

test(test_streaming.py): add unit testing for custom stream wrapper

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-26 08:57:44 -07:00
parent 6f11f300fc
commit 2dd2b8a8e3
2 changed files with 87 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

24
litellm/utils.py
View file

@ -422,8 +422,11 @@ class StreamingChoices(OpenAIObject):
else:
self.finish_reason = None
self.index = index
if delta:
self.delta = delta
if delta is not None:
if isinstance(delta, Delta):
self.delta = delta
if isinstance(delta, dict):
self.delta = Delta(**delta)
else:
self.delta = Delta()
if enhancements is not None:
@ -491,14 +494,25 @@ class ModelResponse(OpenAIObject):
):
if stream is not None and stream == True:
object = "chat.completion.chunk"
choices = [StreamingChoices()]
if choices is not None and isinstance(choices, list):
new_choices = []
for choice in choices:
_new_choice = StreamingChoices(**choice)
new_choices.append(_new_choice)
choices = new_choices
else:
choices = [StreamingChoices()]
else:
if model in litellm.open_ai_embedding_models:
object = "embedding"
else:
object = "chat.completion"
if choices:
choices = [Choices(*choices)]
if choices is not None and isinstance(choices, list):
new_choices = []
for choice in choices:
_new_choice = Choices(**choice)
new_choices.append(_new_choice)
choices = new_choices
else:
choices = [Choices()]
if id is None: