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(main.py): Handle bedrock tool calling in stream_chunk_builder

Browse source 
Fixes #5022.

The streaming chunks from Anthropic seem to violate an assumption
that is implicit in the stream_chunk_builder implementation: that
only tool_calls OR function_calls OR content will appear in a
streamed response. The repro in #5022 shows that you can get
content followed by tool calls.

These changes properly handle these combinations by building
separate lists of each type of chunk (note that in theory a chunk
could appear in multiple lists, e.g. both delta.tool_calls and
delta.content being present on one chunk).
This commit is contained in:
Joe Cheng 2024-08-02 12:41:13 -07:00
parent 52057b2482
commit 0f301a120f
1 changed files with 35 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

53
litellm/main.py
View file

@ -5078,12 +5078,16 @@ def stream_chunk_builder(
combined_content = "" combined_content = ""
combined_arguments = "" combined_arguments = ""
if ( tool_call_chunks = [
"tool_calls" in chunks[0]["choices"][0]["delta"] chunk
and chunks[0]["choices"][0]["delta"]["tool_calls"] is not None for chunk in chunks
): if "tool_calls" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["tool_calls"] is not None
]
if len(tool_call_chunks) > 0:
argument_list = [] argument_list = []
delta = chunks[0]["choices"][0]["delta"] delta = tool_call_chunks[0]["choices"][0]["delta"]
message = response["choices"][0]["message"] message = response["choices"][0]["message"]
message["tool_calls"] = [] message["tool_calls"] = []
id = None id = None
@ -5094,7 +5098,7 @@ def stream_chunk_builder(
prev_id = None prev_id = None
curr_id = None curr_id = None
curr_index = 0 curr_index = 0
for chunk in chunks: for chunk in tool_call_chunks:
choices = chunk["choices"] choices = chunk["choices"]
for choice in choices: for choice in choices:
delta = choice.get("delta", {}) delta = choice.get("delta", {})
@ -5140,12 +5144,17 @@ def stream_chunk_builder(
) )
response["choices"][0]["message"]["content"] = None response["choices"][0]["message"]["content"] = None
response["choices"][0]["message"]["tool_calls"] = tool_calls_list response["choices"][0]["message"]["tool_calls"] = tool_calls_list
elif (
"function_call" in chunks[0]["choices"][0]["delta"] function_call_chunks = [
and chunks[0]["choices"][0]["delta"]["function_call"] is not None chunk
): for chunk in chunks
if "function_calls" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["function_calls"] is not None
]
if len(function_call_chunks) > 0:
argument_list = [] argument_list = []
delta = chunks[0]["choices"][0]["delta"] delta = function_call_chunks[0]["choices"][0]["delta"]
function_call = delta.get("function_call", "") function_call = delta.get("function_call", "")
function_call_name = function_call.name function_call_name = function_call.name
@ -5153,7 +5162,7 @@ def stream_chunk_builder(
message["function_call"] = {} message["function_call"] = {}
message["function_call"]["name"] = function_call_name message["function_call"]["name"] = function_call_name
for chunk in chunks: for chunk in function_call_chunks:
choices = chunk["choices"] choices = chunk["choices"]
for choice in choices: for choice in choices:
delta = choice.get("delta", {}) delta = choice.get("delta", {})
@ -5170,7 +5179,15 @@ def stream_chunk_builder(
response["choices"][0]["message"]["function_call"][ response["choices"][0]["message"]["function_call"][
"arguments" "arguments"
] = combined_arguments ] = combined_arguments
else:
content_chunks = [
chunk
for chunk in chunks
if "content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["content"] is not None
]
if len(content_chunks) > 0:
for chunk in chunks: for chunk in chunks:
choices = chunk["choices"] choices = chunk["choices"]
for choice in choices: for choice in choices:
@ -5186,12 +5203,12 @@ def stream_chunk_builder(
# Update the "content" field within the response dictionary # Update the "content" field within the response dictionary
response["choices"][0]["message"]["content"] = combined_content response["choices"][0]["message"]["content"] = combined_content
completion_output = ""
if len(combined_content) > 0: if len(combined_content) > 0:
completion_output = combined_content completion_output += combined_content
elif len(combined_arguments) > 0: if len(combined_arguments) > 0:
completion_output = combined_arguments completion_output += combined_arguments
else:
completion_output = ""
# # Update usage information if needed # # Update usage information if needed
prompt_tokens = 0 prompt_tokens = 0
completion_tokens = 0 completion_tokens = 0