AVB/bedrock-access-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support multiple tool calls

Browse Source
This commit is contained in:
Aiden Dai
2024-06-11 16:58:26 +08:00
parent 56786f9e32
commit b3509ee0f0
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/api/models/bedrock.py
View File

@@ -398,10 +398,11 @@ class BedrockModel(BaseChatModel):
)
if finish_reason == "tool_use":
# https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html#tool-use-examples
tool_calls = []
for part in content:
if "toolUse" in part:
tool = part["toolUse"]
message.tool_calls = [
tool_calls.append(
ToolCall(
id=tool["toolUseId"],
type="function",
@@ -410,7 +411,8 @@ class BedrockModel(BaseChatModel):
arguments=json.dumps(tool["input"]),
),
)
]
)
message.tool_calls = tool_calls
message.content = None
else:
message.content = content[0]["text"]
@@ -450,7 +452,6 @@ class BedrockModel(BaseChatModel):
finish_reason = None
message = None
usage = None
if "messageStart" in chunk:
message = ChatResponseMessage(
role=chunk["messageStart"]["role"],
@@ -460,10 +461,12 @@ class BedrockModel(BaseChatModel):
# tool call start
delta = chunk["contentBlockStart"]["start"]
if "toolUse" in delta:
# first index is content
index = chunk["contentBlockStart"]["contentBlockIndex"] - 1
message = ChatResponseMessage(
tool_calls=[
ToolCall(
index=0,
index=index,
type="function",
id=delta["toolUse"]["toolUseId"],
function=ResponseFunction(
@@ -482,10 +485,11 @@ class BedrockModel(BaseChatModel):
)
else:
# tool use
index = chunk["contentBlockDelta"]["contentBlockIndex"] - 1
message = ChatResponseMessage(
tool_calls=[
ToolCall(
index=0,
index=index,
function=ResponseFunction(
arguments=delta["toolUse"]["input"],
)