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

v0 sagemaker_stream

Browse source
This commit is contained in:
ishaan-jaff 2024-01-22 21:53:16 -08:00
parent 9652280c15
commit c8084bb9d9
1 changed files with 2 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

13
litellm/utils.py
View file

@ -7715,18 +7715,9 @@ class CustomStreamWrapper:
self.sent_last_chunk = True
elif self.custom_llm_provider == "sagemaker":
print_verbose(f"ENTERS SAGEMAKER STREAMING")
if len(self.completion_stream) == 0:
if self.sent_last_chunk:
raise StopIteration
else:
model_response.choices[0].finish_reason = "stop"
self.sent_last_chunk = True
new_chunk = self.completion_stream
print_verbose(f"sagemaker chunk: {new_chunk}")
new_chunk = next(self.completion_stream)
completion_obj["content"] = new_chunk
self.completion_stream = self.completion_stream[
len(self.completion_stream) :
]
elif self.custom_llm_provider == "petals":
if len(self.completion_stream) == 0:
if self.sent_last_chunk: