From 90dd60fa711a8fb6b7e531d7555bfb51ff7ceff7 Mon Sep 17 00:00:00 2001 From: Joe Cheng Date: Fri, 2 Aug 2024 12:41:13 -0700 Subject: [PATCH 1/4] fix(main.py): Handle bedrock tool calling in stream_chunk_builder 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). --- litellm/main.py | 53 ++++++++++++++++++++++++++++++++----------------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 67b935a55c..989e0b1106 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5078,12 +5078,16 @@ def stream_chunk_builder( combined_content = "" combined_arguments = "" - if ( - "tool_calls" in chunks[0]["choices"][0]["delta"] - and chunks[0]["choices"][0]["delta"]["tool_calls"] is not None - ): + tool_call_chunks = [ + chunk + 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 = [] - delta = chunks[0]["choices"][0]["delta"] + delta = tool_call_chunks[0]["choices"][0]["delta"] message = response["choices"][0]["message"] message["tool_calls"] = [] id = None @@ -5094,7 +5098,7 @@ def stream_chunk_builder( prev_id = None curr_id = None curr_index = 0 - for chunk in chunks: + for chunk in tool_call_chunks: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) @@ -5140,12 +5144,17 @@ def stream_chunk_builder( ) response["choices"][0]["message"]["content"] = None response["choices"][0]["message"]["tool_calls"] = tool_calls_list - elif ( - "function_call" in chunks[0]["choices"][0]["delta"] - and chunks[0]["choices"][0]["delta"]["function_call"] is not None - ): + + function_call_chunks = [ + 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 = [] - delta = chunks[0]["choices"][0]["delta"] + delta = function_call_chunks[0]["choices"][0]["delta"] function_call = delta.get("function_call", "") function_call_name = function_call.name @@ -5153,7 +5162,7 @@ def stream_chunk_builder( message["function_call"] = {} message["function_call"]["name"] = function_call_name - for chunk in chunks: + for chunk in function_call_chunks: choices = chunk["choices"] for choice in choices: delta = choice.get("delta", {}) @@ -5170,7 +5179,15 @@ def stream_chunk_builder( response["choices"][0]["message"]["function_call"][ "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: choices = chunk["choices"] for choice in choices: @@ -5186,12 +5203,12 @@ def stream_chunk_builder( # Update the "content" field within the response dictionary response["choices"][0]["message"]["content"] = combined_content + completion_output = "" if len(combined_content) > 0: - completion_output = combined_content - elif len(combined_arguments) > 0: - completion_output = combined_arguments - else: - completion_output = "" + completion_output += combined_content + if len(combined_arguments) > 0: + completion_output += combined_arguments + # # Update usage information if needed prompt_tokens = 0 completion_tokens = 0 From 33f4411f17b73e62332ff658ebadcb9579b99783 Mon Sep 17 00:00:00 2001 From: Joe Cheng Date: Fri, 2 Aug 2024 13:05:23 -0700 Subject: [PATCH 2/4] Fix tool call coalescing The previous code seemed to assume that the tool call index property started at 0, but Anthropic sometimes returns them starting at 1. This was causing an extra null-ish tool call to be materialized. --- litellm/main.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 989e0b1106..36267aec1d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5094,7 +5094,7 @@ def stream_chunk_builder( name = None type = None tool_calls_list = [] - prev_index = 0 + prev_index = None prev_id = None curr_id = None curr_index = 0 @@ -5120,6 +5120,8 @@ def stream_chunk_builder( name = tool_calls[0].function.name if tool_calls[0].type: type = tool_calls[0].type + if prev_index is None: + prev_index = curr_index if curr_index != prev_index: # new tool call combined_arguments = "".join(argument_list) tool_calls_list.append( @@ -5138,6 +5140,7 @@ def stream_chunk_builder( tool_calls_list.append( { "id": id, + "index": curr_index, "function": {"arguments": combined_arguments, "name": name}, "type": type, } From 1fbfc09b443337b531be863b8a631b8c3bd8e136 Mon Sep 17 00:00:00 2001 From: Joe Cheng Date: Fri, 2 Aug 2024 20:51:03 -0700 Subject: [PATCH 3/4] Add unit test --- litellm/tests/stream_chunk_testdata.py | 543 +++++++++++++++++++++ litellm/tests/test_stream_chunk_builder.py | 23 + 2 files changed, 566 insertions(+) create mode 100644 litellm/tests/stream_chunk_testdata.py diff --git a/litellm/tests/stream_chunk_testdata.py b/litellm/tests/stream_chunk_testdata.py new file mode 100644 index 0000000000..6be9d1ebdf --- /dev/null +++ b/litellm/tests/stream_chunk_testdata.py @@ -0,0 +1,543 @@ +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponse, + StreamingChoices, +) + +chunks = [ + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="To answer", + role="assistant", + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" your", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" question about", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" how", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" many rows are in the ", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="'users' table, I", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="'ll", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" need to", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" run", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" a SQL query.", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" Let", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" me", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" ", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="do that for", + role=None, + function_call=None, + tool_calls=None, + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content=" you.", role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="toolu_01H3AjkLpRtGQrof13CBnWfK", + function=Function(arguments="", name="sql_query"), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments="", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656356, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='{"', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='query": ', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='"SELECT C', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments="OUNT(*", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments=") ", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments="FROM use", name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id=None, + function=Function(arguments='rs;"}', name=None), + type="function", + index=1, + ) + ], + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), + ModelResponse( + id="chatcmpl-634a6ad3-483a-44a1-8cdd-3befbeb4ac2f", + choices=[ + StreamingChoices( + finish_reason="tool_calls", + index=0, + delta=Delta( + content=None, role=None, function_call=None, tool_calls=None + ), + logprobs=None, + ) + ], + created=1722656357, + model="claude-3-5-sonnet-20240620", + object="chat.completion.chunk", + system_fingerprint=None, + ), +] diff --git a/litellm/tests/test_stream_chunk_builder.py b/litellm/tests/test_stream_chunk_builder.py index 342b070ae7..78d2617f1e 100644 --- a/litellm/tests/test_stream_chunk_builder.py +++ b/litellm/tests/test_stream_chunk_builder.py @@ -18,6 +18,8 @@ from openai import OpenAI import litellm from litellm import completion, stream_chunk_builder +import litellm.tests.stream_chunk_testdata + dotenv.load_dotenv() user_message = "What is the current weather in Boston?" @@ -196,3 +198,24 @@ def test_stream_chunk_builder_litellm_usage_chunks(): # assert prompt tokens are the same assert gemini_pt == stream_rebuilt_pt + + +def test_stream_chunk_builder_litellm_mixed_calls(): + response = stream_chunk_builder(litellm.tests.stream_chunk_testdata.chunks) + assert ( + response.choices[0].message.content + == "To answer your question about how many rows are in the 'users' table, I'll need to run a SQL query. Let me do that for you." + ) + + print(response.choices[0].message.tool_calls[0].to_dict()) + + assert len(response.choices[0].message.tool_calls) == 1 + assert response.choices[0].message.tool_calls[0].to_dict() == { + "index": 1, + "function": { + "arguments": '{"query": "SELECT COUNT(*) FROM users;"}', + "name": "sql_query", + }, + "id": "toolu_01H3AjkLpRtGQrof13CBnWfK", + "type": "function", + } From b7be609d6e3d197e0dec4083bf1ed5d9fa148a42 Mon Sep 17 00:00:00 2001 From: Joe Cheng Date: Sat, 3 Aug 2024 11:58:46 -0700 Subject: [PATCH 4/4] Use correct key name --- litellm/main.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 36267aec1d..f6a43023a3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5151,8 +5151,8 @@ def stream_chunk_builder( function_call_chunks = [ chunk for chunk in chunks - if "function_calls" in chunk["choices"][0]["delta"] - and chunk["choices"][0]["delta"]["function_calls"] is not None + if "function_call" in chunk["choices"][0]["delta"] + and chunk["choices"][0]["delta"]["function_call"] is not None ] if len(function_call_chunks) > 0: