diff --git a/pydantic_ai_slim/pydantic_ai/models/bedrock.py b/pydantic_ai_slim/pydantic_ai/models/bedrock.py index 9fb1524449..106c88c2da 100644 --- a/pydantic_ai_slim/pydantic_ai/models/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/models/bedrock.py @@ -59,6 +59,8 @@ MessageUnionTypeDef, PerformanceConfigurationTypeDef, PromptVariableValuesTypeDef, + ReasoningContentBlockOutputTypeDef, + ReasoningTextBlockTypeDef, SystemContentBlockTypeDef, ToolChoiceTypeDef, ToolConfigurationTypeDef, @@ -276,9 +278,10 @@ async def _process_response(self, response: ConverseResponseTypeDef) -> ModelRes if reasoning_content := item.get('reasoningContent'): reasoning_text = reasoning_content.get('reasoningText') if reasoning_text: # pragma: no branch - thinking_part = ThinkingPart(content=reasoning_text['text']) - if reasoning_signature := reasoning_text.get('signature'): - thinking_part.signature = reasoning_signature + thinking_part = ThinkingPart( + content=reasoning_text['text'], + signature=reasoning_text.get('signature'), + ) items.append(thinking_part) if text := item.get('text'): items.append(TextPart(content=text)) @@ -462,8 +465,19 @@ async def _map_messages( # noqa: C901 if isinstance(item, TextPart): content.append({'text': item.content}) elif isinstance(item, ThinkingPart): - # NOTE: We don't pass the thinking part to Bedrock since it raises an error. - pass + if BedrockModelProfile.from_profile(self.profile).bedrock_send_back_thinking_parts: + reasoning_text: ReasoningTextBlockTypeDef = { + 'text': item.content, + } + if item.signature: + reasoning_text['signature'] = item.signature + reasoning_content: ReasoningContentBlockOutputTypeDef = { + 'reasoningText': reasoning_text, + } + content.append({'reasoningContent': reasoning_content}) + else: + # NOTE: We don't pass the thinking part to Bedrock for models other than Claude since it raises an error. + pass else: assert isinstance(item, ToolCallPart) content.append(self._map_tool_call(item)) @@ -610,7 +624,11 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]: delta = chunk['contentBlockDelta']['delta'] if 'reasoningContent' in delta: if text := delta['reasoningContent'].get('text'): - yield self._parts_manager.handle_thinking_delta(vendor_part_id=index, content=text) + yield self._parts_manager.handle_thinking_delta( + vendor_part_id=index, + content=text, + signature=delta['reasoningContent'].get('signature'), + ) else: # pragma: no cover warnings.warn( f'Only text reasoning content is supported yet, but you got {delta["reasoningContent"]}. ' diff --git a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py index a09aa822de..cf19ce290e 100644 --- a/pydantic_ai_slim/pydantic_ai/providers/bedrock.py +++ b/pydantic_ai_slim/pydantic_ai/providers/bedrock.py @@ -36,6 +36,7 @@ class BedrockModelProfile(ModelProfile): bedrock_supports_tool_choice: bool = True bedrock_tool_result_format: Literal['text', 'json'] = 'text' + bedrock_send_back_thinking_parts: bool = False class BedrockProvider(Provider[BaseClient]): @@ -55,9 +56,9 @@ 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).update( - anthropic_model_profile(model_name) - ), + 'anthropic': lambda model_name: BedrockModelProfile( + bedrock_supports_tool_choice=False, 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) ), diff --git a/tests/models/cassettes/test_bedrock/test_bedrock_anthropic_tool_with_thinking.yaml b/tests/models/cassettes/test_bedrock/test_bedrock_anthropic_tool_with_thinking.yaml new file mode 100644 index 0000000000..c13d5d3fcd --- /dev/null +++ b/tests/models/cassettes/test_bedrock/test_bedrock_anthropic_tool_with_thinking.yaml @@ -0,0 +1,116 @@ +interactions: +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "What is the largest city in the user country?"}]}], "system": + [], "inferenceConfig": {}, "toolConfig": {"tools": [{"toolSpec": {"name": "get_user_country", "inputSchema": {"json": + {"additionalProperties": false, "properties": {}, "type": "object"}}}}]}, "additionalModelRequestFields": {"thinking": + {"type": "enabled", "budget_tokens": 1024}}}' + headers: + amz-sdk-invocation-id: + - !!binary | + MmM5YzRmZDctMmNlZS00Yzk2LWIwZWMtZjMxN2NkZDEwYmM5 + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + content-length: + - '396' + content-type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse + response: + headers: + connection: + - keep-alive + content-length: + - '1085' + content-type: + - application/json + parsed_body: + metrics: + latencyMs: 3896 + output: + message: + content: + - reasoningContent: + reasoningText: + signature: ErcBCkgIBhABGAIiQDYN+P1S3ACL3r24cAMKCrRNiuPPxvmT2uzREPLRyKEUXagRbXn97QCke6L7OEZvlh7NdA/MQNTwMZV8TuB4qPASDLNxYxDx1S3luCIfARoMIwLZXwhsvjjTN72XIjAJrEl5ryAvv6C1+6YMCPC73ffE+kgwB96IcZaOuDFQtyaoWwcFcDPBguM6YNp5e3cqHbDQ3QF5dR4PP5q+3K23pual3pUdT/0e7khyIxXkGAI= + text: |- + The user is asking for the largest city in their country. To answer this, I first need to determine what country the user is from. I can use the `get_user_country` function to retrieve this information. + + Once I have the user's country, I can then provide information about the largest city in that country. + - text: I'll need to check what country you're from to answer that question. + - toolUse: + input: {} + name: get_user_country + toolUseId: tooluse_W9DaUFg4Tj2cRPpndqxWSg + role: assistant + stopReason: tool_use + usage: + cacheReadInputTokenCount: 0 + cacheReadInputTokens: 0 + cacheWriteInputTokenCount: 0 + cacheWriteInputTokens: 0 + inputTokens: 397 + outputTokens: 130 + totalTokens: 527 + status: + code: 200 + message: OK +- request: + body: '{"messages": [{"role": "user", "content": [{"text": "What is the largest city in the user country?"}]}, {"role": + "assistant", "content": [{"reasoningContent": {"reasoningText": {"text": "The user is asking for the largest city in + their country. To answer this, I first need to determine what country the user is from. I can use the `get_user_country` + function to retrieve this information.\n\nOnce I have the user''s country, I can then provide information about the + largest city in that country.", "signature": "ErcBCkgIBhABGAIiQDYN+P1S3ACL3r24cAMKCrRNiuPPxvmT2uzREPLRyKEUXagRbXn97QCke6L7OEZvlh7NdA/MQNTwMZV8TuB4qPASDLNxYxDx1S3luCIfARoMIwLZXwhsvjjTN72XIjAJrEl5ryAvv6C1+6YMCPC73ffE+kgwB96IcZaOuDFQtyaoWwcFcDPBguM6YNp5e3cqHbDQ3QF5dR4PP5q+3K23pual3pUdT/0e7khyIxXkGAI="}}}, + {"text": "I''ll need to check what country you''re from to answer that question."}, {"toolUse": {"toolUseId": "tooluse_W9DaUFg4Tj2cRPpndqxWSg", + "name": "get_user_country", "input": {}}}]}, {"role": "user", "content": [{"toolResult": {"toolUseId": "tooluse_W9DaUFg4Tj2cRPpndqxWSg", + "content": [{"text": "Mexico"}], "status": "success"}}]}], "system": [], "inferenceConfig": {}, "toolConfig": {"tools": + [{"toolSpec": {"name": "get_user_country", "inputSchema": {"json": {"additionalProperties": false, "properties": {}, + "type": "object"}}}}]}, "additionalModelRequestFields": {"thinking": {"type": "enabled", "budget_tokens": 1024}}}' + headers: + amz-sdk-invocation-id: + - !!binary | + ZWM5NzBkMzYtZTZhYi00MjdlLWFmMzItMTBhNTc2ZjBiMWNl + amz-sdk-request: + - !!binary | + YXR0ZW1wdD0x + content-length: + - '1399' + content-type: + - !!binary | + YXBwbGljYXRpb24vanNvbg== + method: POST + uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse + response: + headers: + connection: + - keep-alive + content-length: + - '756' + content-type: + - application/json + parsed_body: + metrics: + latencyMs: 4529 + output: + message: + content: + - text: |- + Based on your location in Mexico, the largest city is Mexico City (Ciudad de México). It's not only the capital but also the most populous city in Mexico with a metropolitan area population of over 21 million people, making it one of the largest urban agglomerations in the world. + + Mexico City is an important cultural, financial, and political center for the country and has a rich history dating back to the Aztec empire when it was known as Tenochtitlán. + role: assistant + stopReason: end_turn + usage: + cacheReadInputTokenCount: 0 + cacheReadInputTokens: 0 + cacheWriteInputTokenCount: 0 + cacheWriteInputTokens: 0 + inputTokens: 539 + outputTokens: 106 + totalTokens: 645 + status: + code: 200 + message: OK +version: 1 diff --git a/tests/models/test_bedrock.py b/tests/models/test_bedrock.py index fad3530758..807bb4e1e1 100644 --- a/tests/models/test_bedrock.py +++ b/tests/models/test_bedrock.py @@ -654,6 +654,29 @@ async def test_bedrock_model_thinking_part(allow_model_requests: None, bedrock_p ) +async def test_bedrock_anthropic_tool_with_thinking(allow_model_requests: None, bedrock_provider: BedrockProvider): + """When using thinking with tool calls in Anthropic, we need to send the thinking part back to the provider. + + This tests the issue raised in https://github.com/pydantic/pydantic-ai/issues/2453. + """ + m = BedrockConverseModel('us.anthropic.claude-3-7-sonnet-20250219-v1:0', provider=bedrock_provider) + settings = BedrockModelSettings( + bedrock_additional_model_requests_fields={'thinking': {'type': 'enabled', 'budget_tokens': 1024}}, + ) + agent = Agent(m, model_settings=settings) + + @agent.tool_plain + async def get_user_country() -> str: + return 'Mexico' + + result = await agent.run('What is the largest city in the user country?') + assert result.output == snapshot("""\ +Based on your location in Mexico, the largest city is Mexico City (Ciudad de México). It's not only the capital but also the most populous city in Mexico with a metropolitan area population of over 21 million people, making it one of the largest urban agglomerations in the world. + +Mexico City is an important cultural, financial, and political center for the country and has a rich history dating back to the Aztec empire when it was known as Tenochtitlán.\ +""") + + async def test_bedrock_group_consecutive_tool_return_parts(bedrock_provider: BedrockProvider): """ Test that consecutive ToolReturnPart objects are grouped into a single user message for Bedrock.