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

fix streaming format

Browse source
This commit is contained in:
ishaan-jaff 2023-08-09 13:41:11 -07:00
parent 731b4fd985
commit c83ccb66ca
2 changed files with 8 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/tests/test_completion.py
View file

@ -164,8 +164,8 @@ def test_completion_replicate_stability_stream():
try: try:
response = completion(model=model_name, messages=messages, stream=True, replicate=True) response = completion(model=model_name, messages=messages, stream=True, replicate=True)
# Add any assertions here to check the response # Add any assertions here to check the response
for result in response: for chunk in response:
print(result) print(chunk['choices'][0]['delta'])
print(response) print(response)
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")

9
litellm/utils.py
View file

@ -625,13 +625,16 @@ class CustomStreamWrapper:
return self return self
def __next__(self): def __next__(self):
completion_obj ={ "role": "assistant", "content": ""}
if self.model in litellm.anthropic_models: if self.model in litellm.anthropic_models:
chunk = next(self.completion_stream) chunk = next(self.completion_stream)
return {"choices": [{"delta": chunk.completion}]} completion_obj["content"] = chunk.completion
elif self.model == "replicate": elif self.model == "replicate":
chunk = next(self.completion_stream) chunk = next(self.completion_stream)
return {"choices": [{"delta": chunk}]} completion_obj["content"] = chunk
elif self.model in litellm.cohere_models: elif self.model in litellm.cohere_models:
chunk = next(self.completion_stream) chunk = next(self.completion_stream)
return {"choices": [{"delta": chunk.text}]} completion_obj["content"] = chunk.text
# return this for all models
return {"choices": [{"delta": completion_obj}]}