From f29de0024a1acda81901938c72131350aaa1f345 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 22 Jan 2024 21:50:40 -0800 Subject: [PATCH] (v0) sagemaker streaming --- litellm/main.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 2d8f2c0c912..a7964e4a42e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1519,10 +1519,12 @@ def completion( # fake streaming for sagemaker print_verbose(f"ENTERS SAGEMAKER CUSTOMSTREAMWRAPPER") - resp_string = model_response["choices"][0]["message"]["content"] + from .llms.sagemaker import TokenIterator + + tokenIterator = TokenIterator(model_response) response = CustomStreamWrapper( - resp_string, - model, + completion_stream=tokenIterator, + model=model, custom_llm_provider="sagemaker", logging_obj=logging, )