From 2f815705ca565a9f78be6d8d0d965dd91c65ab12 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Feb 2024 22:13:35 -0800 Subject: [PATCH] fix(sagemaker.py): use __anext__ --- litellm/llms/sagemaker.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index e91913584..b8ebc05bb 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -85,7 +85,7 @@ class TokenIterator: response_obj["is_finished"] = True response_obj["text"] = line_data["token"]["text"] return response_obj - chunk = await anext(self.byte_iterator) + chunk = await self.byte_iterator.__anext__() self.buffer.seek(0, io.SEEK_END) self.buffer.write(chunk["PayloadPart"]["Bytes"]) except StopAsyncIteration as e: @@ -368,10 +368,8 @@ async def async_streaming( 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(