Send ThinkingParts back to Anthropic used through Bedrock

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"]}. '

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)
             ),

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

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.