diff --git a/litellm/main.py b/litellm/main.py index 67b935a55c..f6a43023a3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5078,23 +5078,27 @@ 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 name = None type = None tool_calls_list = [] - prev_index = 0 + prev_index = None 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", {}) @@ -5116,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( @@ -5134,18 +5140,24 @@ def stream_chunk_builder( tool_calls_list.append( { "id": id, + "index": curr_index, "function": {"arguments": combined_arguments, "name": name}, "type": type, } ) 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_call" in chunk["choices"][0]["delta"] + and chunk["choices"][0]["delta"]["function_call"] 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 +5165,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 +5182,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 +5206,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 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", + }