phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(sagemaker.py): use __anext__

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-12 22:13:35 -08:00
parent b1bc30ee16
commit 2f815705ca
1 changed files with 1 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/sagemaker.py
View file

@ -85,7 +85,7 @@ class TokenIterator:
response_obj["is_finished"] = True response_obj["is_finished"] = True
response_obj["text"] = line_data["token"]["text"] response_obj["text"] = line_data["token"]["text"]
return response_obj return response_obj
chunk = await anext(self.byte_iterator) chunk = await self.byte_iterator.__anext__()
self.buffer.seek(0, io.SEEK_END) self.buffer.seek(0, io.SEEK_END)
self.buffer.write(chunk["PayloadPart"]["Bytes"]) self.buffer.write(chunk["PayloadPart"]["Bytes"])
except StopAsyncIteration as e: except StopAsyncIteration as e:
@ -368,10 +368,8 @@ async def async_streaming(
except Exception as e: except Exception as e:
raise SagemakerError(status_code=500, message=f"{str(e)}") raise SagemakerError(status_code=500, message=f"{str(e)}")
response = response["Body"] response = response["Body"]
# filtered_response = TokenIterator(stream=response, acompletion=True)
async for chunk in response: async for chunk in response:
yield chunk yield chunk
# return
async def async_completion( async def async_completion(