Support of tool choice

This commit is contained in:
Aiden Dai
2024-06-21 10:24:11 +08:00
parent b3509ee0f0
commit 49dd6608a0

View File

@@ -381,6 +381,20 @@ class BedrockModel(BaseChatModel):
self._convert_tool_spec(t.function) for t in chat_request.tools self._convert_tool_spec(t.function) for t in chat_request.tools
] ]
} }
if chat_request.tool_choice:
if isinstance(chat_request.tool_choice, str):
# auto (default) is mapped to {"auto" : {}}
# required is mapped to {"any" : {}}
if chat_request.tool_choice == "required":
args["toolConfig"]["toolChoice"] = {"any": {}}
else:
args["toolConfig"]["toolChoice"] = {"auto": {}}
else:
# Specific tool to use
assert "function" in chat_request.tool_choice
args["toolConfig"]["toolChoice"] = {
"tool": {"name": chat_request.tool_choice["function"].get("name", "")}}
return args return args
def _create_response( def _create_response(