From c8084bb9d99e6b642f15b017028b204876539976 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 22 Jan 2024 21:53:16 -0800 Subject: [PATCH] v0 sagemaker_stream --- litellm/utils.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index afa2830c93f..3e8aae9450b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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: