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

(fix) sagemaker streaming support

Browse source
This commit is contained in:
ishaan-jaff 2024-01-23 12:31:16 -08:00
parent 44e213e842
commit e8cd27f2b7
1 changed files with 3 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/utils.py
View file

@ -7732,10 +7732,8 @@ class CustomStreamWrapper:
]
self.sent_last_chunk = True
elif self.custom_llm_provider == "sagemaker":
print_verbose(f"ENTERS SAGEMAKER STREAMING")
new_chunk = next(self.completion_stream)
completion_obj["content"] = new_chunk
print_verbose(f"ENTERS SAGEMAKER STREAMING for chunk {chunk}")
completion_obj["content"] = chunk
elif self.custom_llm_provider == "petals":
if len(self.completion_stream) == 0:
if self.sent_last_chunk:
@ -7854,7 +7852,7 @@ class CustomStreamWrapper:
completion_obj["role"] = "assistant"
self.sent_first_chunk = True
model_response.choices[0].delta = Delta(**completion_obj)
print_verbose(f"model_response: {model_response}")
print_verbose(f"returning model_response: {model_response}")
return model_response
else:
return