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

fix(bedrock_httpx.py): returning correct finish reason on streaming completion

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-10 14:47:49 -07:00
parent 84652dd946
commit 65baa83928
3 changed files with 10 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/llms/bedrock_httpx.py
View file

@ -1898,8 +1898,10 @@ class AWSEventStreamDecoder:
} }
elif "stopReason" in chunk_data: elif "stopReason" in chunk_data:
finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
is_finished = True
elif "usage" in chunk_data: elif "usage" in chunk_data:
usage = ConverseTokenUsageBlock(**chunk_data["usage"]) # type: ignore usage = ConverseTokenUsageBlock(**chunk_data["usage"]) # type: ignore
response = GenericStreamingChunk( response = GenericStreamingChunk(
text=text, text=text,
tool_use=tool_use, tool_use=tool_use,
@ -1924,7 +1926,11 @@ class AWSEventStreamDecoder:
is_finished = True is_finished = True
finish_reason = "stop" finish_reason = "stop"
######## bedrock.anthropic mappings ############### ######## bedrock.anthropic mappings ###############
elif "contentBlockIndex" in chunk_data: elif (
"contentBlockIndex" in chunk_data
or "stopReason" in chunk_data
or "metrics" in chunk_data
):
return self.converse_chunk_parser(chunk_data=chunk_data) return self.converse_chunk_parser(chunk_data=chunk_data)
######## bedrock.mistral mappings ############### ######## bedrock.mistral mappings ###############
elif "outputs" in chunk_data: elif "outputs" in chunk_data:

3
litellm/tests/test_streaming.py
View file

@ -2573,9 +2573,8 @@ def test_streaming_and_function_calling(model):
# Add any assertions here to check the response # Add any assertions here to check the response
for idx, chunk in enumerate(response): for idx, chunk in enumerate(response):
# continue # continue
# print("\n{}\n".format(chunk)) print("\n{}\n".format(chunk))
if idx == 0: if idx == 0:
print(chunk)
assert ( assert (
chunk.choices[0].delta.tool_calls[0].function.arguments is not None chunk.choices[0].delta.tool_calls[0].function.arguments is not None
) )

5
litellm/utils.py
View file

@ -11458,7 +11458,6 @@ class CustomStreamWrapper:
raise StopIteration raise StopIteration
response_obj: GenericStreamingChunk = chunk response_obj: GenericStreamingChunk = chunk
completion_obj["content"] = response_obj["text"] completion_obj["content"] = response_obj["text"]
if response_obj["is_finished"]: if response_obj["is_finished"]:
self.received_finish_reason = response_obj["finish_reason"] self.received_finish_reason = response_obj["finish_reason"]
@ -11591,7 +11590,7 @@ class CustomStreamWrapper:
# for azure, we need to pass the model from the orignal chunk # for azure, we need to pass the model from the orignal chunk
self.model = chunk.model self.model = chunk.model
response_obj = self.handle_openai_chat_completion_chunk(chunk) response_obj = self.handle_openai_chat_completion_chunk(chunk)
if response_obj == None: if response_obj is None:
return return
completion_obj["content"] = response_obj["text"] completion_obj["content"] = response_obj["text"]
print_verbose(f"completion obj content: {completion_obj['content']}") print_verbose(f"completion obj content: {completion_obj['content']}")
@ -11821,7 +11820,7 @@ class CustomStreamWrapper:
else: else:
return return
elif self.received_finish_reason is not None: elif self.received_finish_reason is not None:
if self.sent_last_chunk == True: if self.sent_last_chunk is True:
raise StopIteration raise StopIteration
# flush any remaining holding chunk # flush any remaining holding chunk
if len(self.holding_chunk) > 0: if len(self.holding_chunk) > 0: