fix: anthropic parallel tool call results. (#653)

Co-authored-by: David Montague <[email protected]>

Author: ioga

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -272,64 +272,56 @@ def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageP
         anthropic_messages: list[MessageParam] = []
         for m in messages:
             if isinstance(m, ModelRequest):
-                for part in m.parts:
-                    if isinstance(part, SystemPromptPart):
-                        system_prompt += part.content
-                    elif isinstance(part, UserPromptPart):
-                        anthropic_messages.append(MessageParam(role='user', content=part.content))
-                    elif isinstance(part, ToolReturnPart):
-                        anthropic_messages.append(
-                            MessageParam(
-                                role='user',
-                                content=[
-                                    ToolResultBlockParam(
-                                        tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                        type='tool_result',
-                                        content=part.model_response_str(),
-                                        is_error=False,
-                                    )
-                                ],
-                            )
+                user_content_params: list[ToolResultBlockParam | TextBlockParam] = []
+                for request_part in m.parts:
+                    if isinstance(request_part, SystemPromptPart):
+                        system_prompt += request_part.content
+                    elif isinstance(request_part, UserPromptPart):
+                        text_block_param = TextBlockParam(type='text', text=request_part.content)
+                        user_content_params.append(text_block_param)
+                    elif isinstance(request_part, ToolReturnPart):
+                        tool_result_block_param = ToolResultBlockParam(
+                            tool_use_id=_guard_tool_call_id(t=request_part, model_source='Anthropic'),
+                            type='tool_result',
+                            content=request_part.model_response_str(),
+                            is_error=False,
                         )
-                    elif isinstance(part, RetryPromptPart):
-                        if part.tool_name is None:
-                            anthropic_messages.append(MessageParam(role='user', content=part.model_response()))
+                        user_content_params.append(tool_result_block_param)
+                    elif isinstance(request_part, RetryPromptPart):
+                        if request_part.tool_name is None:
+                            retry_param = TextBlockParam(type='text', text=request_part.model_response())
                         else:
-                            anthropic_messages.append(
-                                MessageParam(
-                                    role='user',
-                                    content=[
-                                        ToolResultBlockParam(
-                                            tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                            type='tool_result',
-                                            content=part.model_response(),
-                                            is_error=True,
-                                        ),
-                                    ],
-                                )
+                            retry_param = ToolResultBlockParam(
+                                tool_use_id=_guard_tool_call_id(t=request_part, model_source='Anthropic'),
+                                type='tool_result',
+                                content=request_part.model_response(),
+                                is_error=True,
                             )
+                        user_content_params.append(retry_param)
+                anthropic_messages.append(
+                    MessageParam(
+                        role='user',
+                        content=user_content_params,
+                    )
+                )
             elif isinstance(m, ModelResponse):
-                content: list[TextBlockParam | ToolUseBlockParam] = []
-                for item in m.parts:
-                    if isinstance(item, TextPart):
-                        content.append(TextBlockParam(text=item.content, type='text'))
+                assistant_content_params: list[TextBlockParam | ToolUseBlockParam] = []
+                for response_part in m.parts:
+                    if isinstance(response_part, TextPart):
+                        assistant_content_params.append(TextBlockParam(text=response_part.content, type='text'))
                     else:
-                        assert isinstance(item, ToolCallPart)
-                        content.append(self._map_tool_call(item))
-                anthropic_messages.append(MessageParam(role='assistant', content=content))
+                        tool_use_block_param = ToolUseBlockParam(
+                            id=_guard_tool_call_id(t=response_part, model_source='Anthropic'),
+                            type='tool_use',
+                            name=response_part.tool_name,
+                            input=response_part.args_as_dict(),
+                        )
+                        assistant_content_params.append(tool_use_block_param)
+                anthropic_messages.append(MessageParam(role='assistant', content=assistant_content_params))
             else:
                 assert_never(m)
         return system_prompt, anthropic_messages
 
-    @staticmethod
-    def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
-        return ToolUseBlockParam(
-            id=_guard_tool_call_id(t=t, model_source='Anthropic'),
-            type='tool_use',
-            name=t.tool_name,
-            input=t.args_as_dict(),
-        )
-
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolParam:
         return {

tests/models/test_anthropic.py

@@ -332,6 +332,36 @@ async def get_location(loc_name: str) -> str:
     )
 
 
+@pytest.mark.skip(reason='Need to finish implementing this test')  # TODO: David to finish implementing this test
+async def test_multiple_parallel_tool_calls():  # pragma: no cover  # Need to finish the test..
+    async def retrieve_entity_info(name: str) -> str:
+        """Get the knowledge about the given entity."""
+        data = {
+            'alice': "alice is bob's wife",
+            'bob': "bob is alice's husband",
+            'charlie': "charlie is alice's son",
+            'daisy': "daisy is bob's daughter and charlie's younger sister",
+        }
+        return data[name.lower()]
+
+    system_prompt = """
+    Use the `retrieve_entity_info` tool to get information about a specific person.
+    If you need to use `retrieve_entity_info` to get information about multiple people, try
+    to call them in parallel as much as possible.
+    Think step by step and then provide a single most probable concise answer.
+    """
+
+    mock_client = MockAnthropic.create_mock([])
+    agent = Agent(
+        AnthropicModel('claude-3-5-haiku-latest', anthropic_client=mock_client),
+        system_prompt=system_prompt,
+        tools=[retrieve_entity_info],
+    )
+
+    _result = agent.run_sync('Alice, Bob, Charlie and Daisy are a family. Who is the youngest?')
+    # assert ...
+
+
 async def test_anthropic_specific_metadata(allow_model_requests: None) -> None:
     c = completion_message([TextBlock(text='world', type='text')], AnthropicUsage(input_tokens=5, output_tokens=10))
     mock_client = MockAnthropic.create_mock(c)