tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-13 20:23:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #5025 from jcheng5/streaming-tools-anthropic

Browse Source 
Handle bedrock tool calling in stream_chunk_builder
This commit is contained in:
Krish Dholakia
2024-08-03 20:31:05 -07:00
committed by GitHub
parent bd227f7b57 b7be609d6e
commit c5fe467cdd
3 changed files with 605 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/main.py
+39 -19
View File
@@ -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
litellm/tests/stream_chunk_testdata.py
+543
View File
@@ -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,
),
]
litellm/tests/test_stream_chunk_builder.py
+23
View File
@@ -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",
}