(v0) sagemaker streaming

2024-01-22 21:50:40 -08:00 · 2024-01-22 21:50:40 -08:00 · f29de0024a
commit f29de0024a
parent 802fd6cebf
1 changed files with 5 additions and 3 deletions
--- 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,
                )