Skip to content

Commit 3c2f1cf

Browse files
authored
Only send tool choice to Bedrock Converse API for Anthropic and Nova models (#2819)
1 parent 55c6814 commit 3c2f1cf

File tree

3 files changed

+44
-9
lines changed

3 files changed

+44
-9
lines changed

pydantic_ai_slim/pydantic_ai/providers/bedrock.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,19 @@ class BedrockModelProfile(ModelProfile):
3535
ALL FIELDS MUST BE `bedrock_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
3636
"""
3737

38-
bedrock_supports_tool_choice: bool = True
38+
bedrock_supports_tool_choice: bool = False
3939
bedrock_tool_result_format: Literal['text', 'json'] = 'text'
4040
bedrock_send_back_thinking_parts: bool = False
4141

4242

43+
def bedrock_amazon_model_profile(model_name: str) -> ModelProfile | None:
44+
"""Get the model profile for an Amazon model used via Bedrock."""
45+
profile = amazon_model_profile(model_name)
46+
if 'nova' in model_name:
47+
return BedrockModelProfile(bedrock_supports_tool_choice=True).update(profile)
48+
return profile
49+
50+
4351
class BedrockProvider(Provider[BaseClient]):
4452
"""Provider for AWS Bedrock."""
4553

@@ -58,13 +66,13 @@ def client(self) -> BaseClient:
5866
def model_profile(self, model_name: str) -> ModelProfile | None:
5967
provider_to_profile: dict[str, Callable[[str], ModelProfile | None]] = {
6068
'anthropic': lambda model_name: BedrockModelProfile(
61-
bedrock_supports_tool_choice=False, bedrock_send_back_thinking_parts=True
69+
bedrock_supports_tool_choice=True, bedrock_send_back_thinking_parts=True
6270
).update(anthropic_model_profile(model_name)),
6371
'mistral': lambda model_name: BedrockModelProfile(bedrock_tool_result_format='json').update(
6472
mistral_model_profile(model_name)
6573
),
6674
'cohere': cohere_model_profile,
67-
'amazon': amazon_model_profile,
75+
'amazon': bedrock_amazon_model_profile,
6876
'meta': meta_model_profile,
6977
'deepseek': deepseek_model_profile,
7078
}

tests/models/test_bedrock.py

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -844,15 +844,15 @@ async def test_bedrock_mistral_tool_result_format(bedrock_provider: BedrockProvi
844844
)
845845

846846

847-
async def test_bedrock_anthropic_no_tool_choice(bedrock_provider: BedrockProvider):
847+
async def test_bedrock_no_tool_choice(bedrock_provider: BedrockProvider):
848848
my_tool = ToolDefinition(
849849
name='my_tool',
850850
description='This is my tool',
851851
parameters_json_schema={'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}},
852852
)
853853
mrp = ModelRequestParameters(output_mode='tool', function_tools=[my_tool], allow_text_output=False, output_tools=[])
854854

855-
# Models other than Anthropic support tool_choice
855+
# Amazon Nova supports tool_choice
856856
model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider)
857857
tool_config = model._map_tool_config(mrp) # type: ignore[reportPrivateUsage]
858858

@@ -873,10 +873,31 @@ async def test_bedrock_anthropic_no_tool_choice(bedrock_provider: BedrockProvide
873873
}
874874
)
875875

876-
# Anthropic models don't support tool_choice
876+
# Anthropic supports tool_choice
877877
model = BedrockConverseModel('us.anthropic.claude-3-7-sonnet-20250219-v1:0', provider=bedrock_provider)
878878
tool_config = model._map_tool_config(mrp) # type: ignore[reportPrivateUsage]
879879

880+
assert tool_config == snapshot(
881+
{
882+
'tools': [
883+
{
884+
'toolSpec': {
885+
'name': 'my_tool',
886+
'description': 'This is my tool',
887+
'inputSchema': {
888+
'json': {'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}}
889+
},
890+
}
891+
}
892+
],
893+
'toolChoice': {'any': {}},
894+
}
895+
)
896+
897+
# Other models don't support tool_choice
898+
model = BedrockConverseModel('us.meta.llama4-maverick-17b-instruct-v1:0', provider=bedrock_provider)
899+
tool_config = model._map_tool_config(mrp) # type: ignore[reportPrivateUsage]
900+
880901
assert tool_config == snapshot(
881902
{
882903
'tools': [

tests/providers/test_bedrock.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,12 +58,12 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture):
5858
anthropic_profile = provider.model_profile('us.anthropic.claude-3-5-sonnet-20240620-v1:0')
5959
anthropic_model_profile_mock.assert_called_with('claude-3-5-sonnet-20240620')
6060
assert isinstance(anthropic_profile, BedrockModelProfile)
61-
assert not anthropic_profile.bedrock_supports_tool_choice
61+
assert anthropic_profile.bedrock_supports_tool_choice is True
6262

6363
anthropic_profile = provider.model_profile('anthropic.claude-instant-v1')
6464
anthropic_model_profile_mock.assert_called_with('claude-instant')
6565
assert isinstance(anthropic_profile, BedrockModelProfile)
66-
assert not anthropic_profile.bedrock_supports_tool_choice
66+
assert anthropic_profile.bedrock_supports_tool_choice is True
6767

6868
mistral_profile = provider.model_profile('mistral.mistral-large-2407-v1:0')
6969
mistral_model_profile_mock.assert_called_with('mistral-large-2407')
@@ -84,7 +84,13 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture):
8484
assert deepseek_profile is not None
8585
assert deepseek_profile.ignore_streamed_leading_whitespace is True
8686

87-
amazon_profile = provider.model_profile('amazon.titan-text-express-v1')
87+
amazon_profile = provider.model_profile('us.amazon.nova-pro-v1:0')
88+
amazon_model_profile_mock.assert_called_with('nova-pro')
89+
assert isinstance(amazon_profile, BedrockModelProfile)
90+
assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer
91+
assert amazon_profile.bedrock_supports_tool_choice is True
92+
93+
amazon_profile = provider.model_profile('us.amazon.titan-text-express-v1:0')
8894
amazon_model_profile_mock.assert_called_with('titan-text-express')
8995
assert amazon_profile is not None
9096
assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer

0 commit comments

Comments
 (0)