From b1bc30ee16584f5b499f1da68b5a67866ea23638 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 12 Feb 2024 21:18:34 -0800
Subject: [PATCH] feat(sagemaker.py): aioboto3 streaming support

---
 litellm/llms/sagemaker.py       | 85 ++++++++++++++++++++++++++++-----
 litellm/main.py                 |  2 +-
 litellm/tests/test_streaming.py |  4 +-
 litellm/utils.py                |  2 +
 4 files changed, 80 insertions(+), 13 deletions(-)

diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py
index cd107df211f..e919135844c 100644
--- a/litellm/llms/sagemaker.py
+++ b/litellm/llms/sagemaker.py
@@ -30,8 +30,11 @@ import json
 
 
 class TokenIterator:
-    def __init__(self, stream):
-        self.byte_iterator = iter(stream)
+    def __init__(self, stream, acompletion: bool):
+        if acompletion == False:
+            self.byte_iterator = iter(stream)
+        elif acompletion == True:
+            self.byte_iterator = stream
         self.buffer = io.BytesIO()
         self.read_pos = 0
         self.end_of_data = False
@@ -64,6 +67,34 @@ class TokenIterator:
                 self.end_of_data = True
                 return "data: [DONE]"
 
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            while True:
+                self.buffer.seek(self.read_pos)
+                line = self.buffer.readline()
+                if line and line[-1] == ord("\n"):
+                    response_obj = {"text": "", "is_finished": False}
+                    self.read_pos += len(line) + 1
+                    full_line = line[:-1].decode("utf-8")
+                    line_data = json.loads(full_line.lstrip("data:").rstrip("/n"))
+                    if line_data.get("generated_text", None) is not None:
+                        self.end_of_data = True
+                        response_obj["is_finished"] = True
+                    response_obj["text"] = line_data["token"]["text"]
+                    return response_obj
+                chunk = await anext(self.byte_iterator)
+                self.buffer.seek(0, io.SEEK_END)
+                self.buffer.write(chunk["PayloadPart"]["Bytes"])
+        except StopAsyncIteration as e:
+            if self.end_of_data == True:
+                raise e  # Re-raise StopIteration
+            else:
+                self.end_of_data = True
+                return "data: [DONE]"
+
 
 class SagemakerConfig:
     """
@@ -197,15 +228,16 @@ def completion(
         data = json.dumps(
             {"inputs": prompt, "parameters": inference_params, "stream": True}
         ).encode("utf-8")
-        ## LOGGING
-        request_str = f"""
-        response = client.invoke_endpoint_with_response_stream(
-            EndpointName={model},
-            ContentType="application/json",
-            Body={data},
-            CustomAttributes="accept_eula=true",
-        )
-        """  # type: ignore
+        if acompletion == True:
+            response = async_streaming(
+                optional_params=optional_params,
+                encoding=encoding,
+                model_response=model_response,
+                model=model,
+                logging_obj=logging_obj,
+                data=data,
+            )
+            return response
         response = client.invoke_endpoint_with_response_stream(
             EndpointName=model,
             ContentType="application/json",
@@ -311,6 +343,37 @@ def completion(
     return model_response
 
 
+async def async_streaming(
+    optional_params,
+    encoding,
+    model_response: ModelResponse,
+    model: str,
+    logging_obj: Any,
+    data,
+):
+    """
+    Use aioboto3
+    """
+    import aioboto3
+
+    session = aioboto3.Session()
+    async with session.client("sagemaker-runtime", region_name="us-west-2") as client:
+        try:
+            response = await client.invoke_endpoint_with_response_stream(
+                EndpointName=model,
+                ContentType="application/json",
+                Body=data,
+                CustomAttributes="accept_eula=true",
+            )
+        except Exception as e:
+            raise SagemakerError(status_code=500, message=f"{str(e)}")
+        response = response["Body"]
+        # filtered_response = TokenIterator(stream=response, acompletion=True)
+        async for chunk in response:
+            yield chunk
+        # return
+
+
 async def async_completion(
     optional_params,
     encoding,
diff --git a/litellm/main.py b/litellm/main.py
index 31dd4324c5e..2d27dc6b9de 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -1562,7 +1562,7 @@ def completion(
                 print_verbose(f"ENTERS SAGEMAKER CUSTOMSTREAMWRAPPER")
                 from .llms.sagemaker import TokenIterator
 
-                tokenIterator = TokenIterator(model_response)
+                tokenIterator = TokenIterator(model_response, acompletion=acompletion)
                 response = CustomStreamWrapper(
                     completion_stream=tokenIterator,
                     model=model,
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 28a9f9902e9..a5497b539cf 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -876,7 +876,6 @@ async def test_sagemaker_streaming_async():
             temperature=0.7,
             stream=True,
         )
-
         # Add any assertions here to check the response
         print(response)
         complete_response = ""
@@ -900,6 +899,9 @@ async def test_sagemaker_streaming_async():
         pytest.fail(f"An exception occurred - {str(e)}")
 
 
+asyncio.run(test_sagemaker_streaming_async())
+
+
 def test_completion_sagemaker_stream():
     try:
         response = completion(
diff --git a/litellm/utils.py b/litellm/utils.py
index ac26e08640d..b67025a49c1 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8691,6 +8691,8 @@ class CustomStreamWrapper:
                 or self.custom_llm_provider == "ollama"
                 or self.custom_llm_provider == "ollama_chat"
                 or self.custom_llm_provider == "vertex_ai"
+                or self.custom_llm_provider == "sagemaker"
+                or self.custom_llm_provider in litellm.openai_compatible_endpoints
             ):
                 print_verbose(
                     f"value of async completion stream: {self.completion_stream}"