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

Merge pull request #1561 from BerriAI/litellm_sagemaker_streaming

Browse source 
[Feat] Add REAL Sagemaker streaming
This commit is contained in:
Ishaan Jaff 2024-01-22 22:10:20 -08:00 • committed by GitHub
commit 6d105754d7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 141 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

13
litellm/utils.py
View file

@ -7736,18 +7736,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: