Merge pull request #1561 from BerriAI/litellm_sagemaker_streaming

[Feat] Add REAL Sagemaker streaming
2025-04-27 11:43:54 +00:00 · 2024-01-22 22:10:20 -08:00 · 2024-01-22 22:10:20 -08:00 · 97dd61a6cb
commit 97dd61a6cb
parent 2af28dbf79 83f8e0c177
5 changed files with 141 additions and 15 deletions
--- a/litellm/main.py
+++ b/litellm/main.py
@ -1520,10 +1520,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,
                )