diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 07ba97ace..b4f3866b8 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -164,8 +164,8 @@ def test_completion_replicate_stability_stream(): try: response = completion(model=model_name, messages=messages, stream=True, replicate=True) # Add any assertions here to check the response - for result in response: - print(result) + for chunk in response: + print(chunk['choices'][0]['delta']) print(response) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/litellm/utils.py b/litellm/utils.py index b81e9bc0d..b961fe812 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -625,13 +625,16 @@ class CustomStreamWrapper: return self def __next__(self): + completion_obj ={ "role": "assistant", "content": ""} if self.model in litellm.anthropic_models: chunk = next(self.completion_stream) - return {"choices": [{"delta": chunk.completion}]} + completion_obj["content"] = chunk.completion elif self.model == "replicate": chunk = next(self.completion_stream) - return {"choices": [{"delta": chunk}]} + completion_obj["content"] = chunk elif self.model in litellm.cohere_models: 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}]}