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

assume index is not always in stream chunk

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-15 19:42:03 -07:00
parent e500e14655
commit 6d9121ea78
1 changed files with 9 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

13
litellm/llms/sagemaker.py
View file

@ -808,16 +808,21 @@ class AWSEventStreamDecoder:
self.model = model self.model = model
self.parser = EventStreamJSONParser() self.parser = EventStreamJSONParser()
self.content_blocks: List = [] self.content_blocks: List = []
self.index = 0
def _chunk_parser(self, chunk_data: dict) -> GChunk: def _chunk_parser(self, chunk_data: dict) -> GChunk:
verbose_logger.debug("in sagemaker chunk parser, chunk_data %s", chunk_data) verbose_logger.debug("in sagemaker chunk parser, chunk_data %s", chunk_data)
_token = chunk_data["token"] _token = chunk_data.get("token", {}) or {}
_index = chunk_data["index"] _index = chunk_data.get("index", None)
if _index is None:
_index = self.index
self.index += 1
is_finished = False is_finished = False
finish_reason = "" finish_reason = ""
if _token["text"] == "<|endoftext|>": _text = _token.get("text", "")
if _text == "<|endoftext|>":
return GChunk( return GChunk(
text="", text="",
index=_index, index=_index,
@ -826,7 +831,7 @@ class AWSEventStreamDecoder:
) )
return GChunk( return GChunk(
text=_token["text"], text=_text,
index=_index, index=_index,
is_finished=is_finished, is_finished=is_finished,
finish_reason=finish_reason, finish_reason=finish_reason,