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

test(test_streaming.py): fix test

Browse source
This commit is contained in:
Krrish Dholakia 2024-04-25 20:22:18 -07:00
parent 5307510592
commit 9eb75cc159
2 changed files with 16 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

14
litellm/tests/test_streaming.py
View file

@ -2723,8 +2723,18 @@ def test_aamazing_unit_test_custom_stream_wrapper_n():
chunk_list = []
for chunk in chunks:
_chunk = litellm.ModelResponse(**chunk, stream=True)
chunk_list.append(_chunk)
new_chunk = litellm.ModelResponse(stream=True, id=chunk["id"])
if "choices" in chunk and isinstance(chunk["choices"], list):
print("INSIDE CHUNK CHOICES!")
new_choices = []
for choice in chunk["choices"]:
if isinstance(choice, litellm.utils.StreamingChoices):
_new_choice = choice
elif isinstance(choice, dict):
_new_choice = litellm.utils.StreamingChoices(**choice)
new_choices.append(_new_choice)
new_chunk.choices = new_choices
chunk_list.append(new_chunk)
completion_stream = ModelResponseListIterator(model_responses=chunk_list)

7
litellm/utils.py
View file

@ -7103,9 +7103,10 @@ def convert_to_model_response_object(
model_response_object.model = response_object["model"]
if start_time is not None and end_time is not None:
model_response_object._response_ms = ( # type: ignore
end_time - start_time
).total_seconds() * 1000
if isinstance(start_time, type(end_time)):
model_response_object._response_ms = ( # type: ignore
end_time - start_time
).total_seconds() * 1000
if hidden_params is not None:
model_response_object._hidden_params = hidden_params