From 179291f34f8f6efb55c9db54f060b18b8074afbe Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 5 Sep 2025 17:01:13 +0000 Subject: [PATCH 1/3] Only send tool choice to Bedrock Converse API for Anthropic and Nova models --- pydantic_ai_slim/pydantic_ai/providers/bedrock.py | 14 +++++++++++--- tests/providers/test_bedrock.py | 11 ++++++----- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index 9a5ce388ea..3e1732fe35 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -35,11 +35,19 @@ class BedrockModelProfile(ModelProfile): ALL FIELDS MUST BE `bedrock_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS. """ - bedrock_supports_tool_choice: bool = True + bedrock_supports_tool_choice: bool = False bedrock_tool_result_format: Literal['text', 'json'] = 'text' bedrock_send_back_thinking_parts: bool = False +def bedrock_amazon_model_profile(model_name: str) -> ModelProfile | None: + """Get the model profile for an Amazon model used via Bedrock.""" + profile = amazon_model_profile(model_name) + if 'nova' in model_name: # pragma: no branch + return BedrockModelProfile(bedrock_supports_tool_choice=True).update(profile) + return profile + + class BedrockProvider(Provider[BaseClient]): """Provider for AWS Bedrock.""" @@ -58,13 +66,13 @@ def client(self) -> BaseClient: def model_profile(self, model_name: str) -> ModelProfile | None: provider_to_profile: dict[str, Callable[[str], ModelProfile | None]] = { 'anthropic': lambda model_name: BedrockModelProfile( - bedrock_supports_tool_choice=False, bedrock_send_back_thinking_parts=True + bedrock_supports_tool_choice=True, bedrock_send_back_thinking_parts=True ).update(anthropic_model_profile(model_name)), 'mistral': lambda model_name: BedrockModelProfile(bedrock_tool_result_format='json').update( mistral_model_profile(model_name) ), 'cohere': cohere_model_profile, - 'amazon': amazon_model_profile, + 'amazon': bedrock_amazon_model_profile, 'meta': meta_model_profile, 'deepseek': deepseek_model_profile, } diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index 4aa7b73058..739c089876 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -58,12 +58,12 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture): anthropic_profile = provider.model_profile('us.anthropic.claude-3-5-sonnet-20240620-v1:0') anthropic_model_profile_mock.assert_called_with('claude-3-5-sonnet-20240620') assert isinstance(anthropic_profile, BedrockModelProfile) - assert not anthropic_profile.bedrock_supports_tool_choice + assert anthropic_profile.bedrock_supports_tool_choice is True anthropic_profile = provider.model_profile('anthropic.claude-instant-v1') anthropic_model_profile_mock.assert_called_with('claude-instant') assert isinstance(anthropic_profile, BedrockModelProfile) - assert not anthropic_profile.bedrock_supports_tool_choice + assert anthropic_profile.bedrock_supports_tool_choice is True mistral_profile = provider.model_profile('mistral.mistral-large-2407-v1:0') mistral_model_profile_mock.assert_called_with('mistral-large-2407') @@ -84,10 +84,11 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture): assert deepseek_profile is not None assert deepseek_profile.ignore_streamed_leading_whitespace is True - amazon_profile = provider.model_profile('amazon.titan-text-express-v1') - amazon_model_profile_mock.assert_called_with('titan-text-express') - assert amazon_profile is not None + amazon_profile = provider.model_profile('us.amazon.nova-pro-v1:0') + amazon_model_profile_mock.assert_called_with('nova-pro') + assert isinstance(amazon_profile, BedrockModelProfile) assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer + assert amazon_profile.bedrock_supports_tool_choice is True unknown_model = provider.model_profile('unknown-model') assert unknown_model is None From d721a7f6ba85b777048efbe281aac4e5c2974559 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 5 Sep 2025 17:12:31 +0000 Subject: [PATCH 2/3] Update test --- tests/models/test_bedrock.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/tests/models/test_bedrock.py b/tests/models/test_bedrock.py index 38db10678e..1e25292e12 100644 --- a/tests/models/test_bedrock.py +++ b/tests/models/test_bedrock.py @@ -844,7 +844,7 @@ async def test_bedrock_mistral_tool_result_format(bedrock_provider: BedrockProvi ) -async def test_bedrock_anthropic_no_tool_choice(bedrock_provider: BedrockProvider): +async def test_bedrock_no_tool_choice(bedrock_provider: BedrockProvider): my_tool = ToolDefinition( name='my_tool', description='This is my tool', @@ -852,7 +852,7 @@ async def test_bedrock_anthropic_no_tool_choice(bedrock_provider: BedrockProvide ) mrp = ModelRequestParameters(output_mode='tool', function_tools=[my_tool], allow_text_output=False, output_tools=[]) - # Models other than Anthropic support tool_choice + # Amazon Nova supports tool_choice model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider) tool_config = model._map_tool_config(mrp) # type: ignore[reportPrivateUsage] @@ -873,10 +873,31 @@ async def test_bedrock_anthropic_no_tool_choice(bedrock_provider: BedrockProvide } ) - # Anthropic models don't support tool_choice + # Anthropic supports tool_choice model = BedrockConverseModel('us.anthropic.claude-3-7-sonnet-20250219-v1:0', provider=bedrock_provider) tool_config = model._map_tool_config(mrp) # type: ignore[reportPrivateUsage] + assert tool_config == snapshot( + { + 'tools': [ + { + 'toolSpec': { + 'name': 'my_tool', + 'description': 'This is my tool', + 'inputSchema': { + 'json': {'type': 'object', 'title': 'Result', 'properties': {'spam': {'type': 'number'}}} + }, + } + } + ], + 'toolChoice': {'any': {}}, + } + ) + + # Other models don't support tool_choice + model = BedrockConverseModel('us.meta.llama4-maverick-17b-instruct-v1:0', provider=bedrock_provider) + tool_config = model._map_tool_config(mrp) # type: ignore[reportPrivateUsage] + assert tool_config == snapshot( { 'tools': [ From b7c7ba9348ffb41495a42f6eab8ec0339b58e017 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 5 Sep 2025 17:43:51 +0000 Subject: [PATCH 3/3] Fix coverage --- pydantic_ai_slim/pydantic_ai/providers/bedrock.py | 2 +- tests/providers/test_bedrock.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index 3e1732fe35..b60c43cbc0 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -43,7 +43,7 @@ class BedrockModelProfile(ModelProfile): def bedrock_amazon_model_profile(model_name: str) -> ModelProfile | None: """Get the model profile for an Amazon model used via Bedrock.""" profile = amazon_model_profile(model_name) - if 'nova' in model_name: # pragma: no branch + if 'nova' in model_name: return BedrockModelProfile(bedrock_supports_tool_choice=True).update(profile) return profile diff --git a/tests/providers/test_bedrock.py b/tests/providers/test_bedrock.py index 739c089876..f026dc3377 100644 --- a/tests/providers/test_bedrock.py +++ b/tests/providers/test_bedrock.py @@ -90,6 +90,11 @@ def test_bedrock_provider_model_profile(env: TestEnv, mocker: MockerFixture): assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer assert amazon_profile.bedrock_supports_tool_choice is True + amazon_profile = provider.model_profile('us.amazon.titan-text-express-v1:0') + amazon_model_profile_mock.assert_called_with('titan-text-express') + assert amazon_profile is not None + assert amazon_profile.json_schema_transformer == InlineDefsJsonSchemaTransformer + unknown_model = provider.model_profile('unknown-model') assert unknown_model is None