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

sagemaker streaming

Browse source
This commit is contained in:
ishaan-jaff 2023-09-21 16:14:44 -07:00
parent a8c0f46111
commit 6add152818
3 changed files with 36 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/utils.py
View file

@ -2720,6 +2720,14 @@ class CustomStreamWrapper:
completion_obj["content"] = self.handle_cohere_chunk(chunk)
elif self.custom_llm_provider == "bedrock":
completion_obj["content"] = self.handle_bedrock_stream()
elif self.custom_llm_provider == "sagemaker":
if len(self.completion_stream)==0:
raise StopIteration
chunk_size = 30
new_chunk = self.completion_stream[:chunk_size]
completion_obj["content"] = new_chunk
self.completion_stream = self.completion_stream[chunk_size:]
time.sleep(0.05)
else: # openai chat/azure models
chunk = next(self.completion_stream)
model_response = chunk