Fix streaming tool calls

Author: DouweM

docs/output.md

@@ -139,7 +139,6 @@ from pydantic import BaseModel
 
 from pydantic_ai import Agent, ModelRetry, RunContext
 from pydantic_ai.exceptions import UnexpectedModelBehavior
-from pydantic_ai.output import ToolRetryError
 
 
 class Row(BaseModel):

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -561,7 +561,7 @@ def multi_modal_content_identifier(identifier: str | bytes) -> str:
     return hashlib.sha1(identifier).hexdigest()[:6]
 
 
-async def process_function_tools(
+async def process_function_tools(  # noqa: C901
     toolset: AbstractToolset[DepsT],
     tool_calls: list[_messages.ToolCallPart],
     final_result: result.FinalResult[NodeRunEndT] | None,
@@ -646,70 +646,72 @@ async def process_function_tools(
 
     user_parts: list[_messages.UserPromptPart] = []
 
-    include_content = (
-        ctx.deps.instrumentation_settings is not None and ctx.deps.instrumentation_settings.include_content
-    )
+    if calls_to_run:
+        include_content = (
+            ctx.deps.instrumentation_settings is not None and ctx.deps.instrumentation_settings.include_content
+        )
 
-    # Run all tool tasks in parallel
-    results_by_index: dict[int, _messages.ModelRequestPart] = {}
-    with ctx.deps.tracer.start_as_current_span(
-        'running tools',
-        attributes={
-            'tools': [call.tool_name for call in calls_to_run],
-            'logfire.msg': f'running {len(calls_to_run)} tool{"" if len(calls_to_run) == 1 else "s"}',
-        },
-    ):
-        tasks = [
-            asyncio.create_task(
-                _call_function_tool(toolset, call, run_context, ctx.deps.tracer, include_content), name=call.tool_name
-            )
-            for call in calls_to_run
-        ]
-
-        pending = tasks
-        while pending:
-            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-            for task in done:
-                index = tasks.index(task)
-                tool_result = task.result()
-                yield _messages.FunctionToolResultEvent(tool_result, tool_call_id=tool_result.tool_call_id)
-
-                if isinstance(tool_result, _messages.RetryPromptPart):
-                    results_by_index[index] = tool_result
-                elif isinstance(tool_result, _messages.ToolReturnPart):
-
-                    def process_content(content: Any) -> Any:
-                        if isinstance(content, _messages.MultiModalContentTypes):
-                            if isinstance(content, _messages.BinaryContent):
-                                identifier = multi_modal_content_identifier(content.data)
+        # Run all tool tasks in parallel
+        results_by_index: dict[int, _messages.ModelRequestPart] = {}
+        with ctx.deps.tracer.start_as_current_span(
+            'running tools',
+            attributes={
+                'tools': [call.tool_name for call in calls_to_run],
+                'logfire.msg': f'running {len(calls_to_run)} tool{"" if len(calls_to_run) == 1 else "s"}',
+            },
+        ):
+            tasks = [
+                asyncio.create_task(
+                    _call_function_tool(toolset, call, run_context, ctx.deps.tracer, include_content),
+                    name=call.tool_name,
+                )
+                for call in calls_to_run
+            ]
+
+            pending = tasks
+            while pending:
+                done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+                for task in done:
+                    index = tasks.index(task)
+                    tool_result = task.result()
+                    yield _messages.FunctionToolResultEvent(tool_result, tool_call_id=tool_result.tool_call_id)
+
+                    if isinstance(tool_result, _messages.RetryPromptPart):
+                        results_by_index[index] = tool_result
+                    elif isinstance(tool_result, _messages.ToolReturnPart):
+
+                        def process_content(content: Any) -> Any:
+                            if isinstance(content, _messages.MultiModalContentTypes):
+                                if isinstance(content, _messages.BinaryContent):
+                                    identifier = multi_modal_content_identifier(content.data)
+                                else:
+                                    identifier = multi_modal_content_identifier(content.url)
+
+                                user_parts.append(
+                                    _messages.UserPromptPart(
+                                        content=[f'This is file {identifier}:', content],
+                                        timestamp=tool_result.timestamp,
+                                        part_kind='user-prompt',
+                                    )
+                                )
+                                return f'See file {identifier}'
                             else:
-                                identifier = multi_modal_content_identifier(content.url)
+                                return content
 
-                            user_parts.append(
-                                _messages.UserPromptPart(
-                                    content=[f'This is file {identifier}:', content],
-                                    timestamp=tool_result.timestamp,
-                                    part_kind='user-prompt',
-                                )
-                            )
-                            return f'See file {identifier}'
+                        if isinstance(tool_result.content, list):
+                            contents = cast(list[Any], tool_result.content)  # type: ignore
+                            tool_result.content = [process_content(content) for content in contents]
                         else:
-                            return content
+                            tool_result.content = process_content(tool_result.content)
 
-                    if isinstance(tool_result.content, list):
-                        contents = cast(list[Any], tool_result.content)  # type: ignore
-                        tool_result.content = [process_content(content) for content in contents]
+                        results_by_index[index] = tool_result
                     else:
-                        tool_result.content = process_content(tool_result.content)
-
-                    results_by_index[index] = tool_result
-                else:
-                    assert_never(tool_result)
+                        assert_never(tool_result)
 
-    # We append the results at the end, rather than as they are received, to retain a consistent ordering
-    # This is mostly just to simplify testing
-    for k in sorted(results_by_index):
-        parts.append(results_by_index[k])
+        # We append the results at the end, rather than as they are received, to retain a consistent ordering
+        # This is mostly just to simplify testing
+        for k in sorted(results_by_index):
+            parts.append(results_by_index[k])
 
     deferred_calls: list[_messages.ToolCallPart] = []
     for call in tool_calls_by_kind['deferred']:
@@ -739,6 +741,7 @@ def process_content(content: Any) -> Any:
     parts.extend(user_parts)
 
     if final_result:
+        # TODO: Use some better "box" object
         final_result_holder.append(final_result)
 
 

tests/models/test_anthropic.py

@@ -1700,7 +1700,7 @@ class CityLocation(BaseModel):
 
     agent = Agent(m, output_type=NativeOutput(CityLocation))
 
-    with pytest.raises(UserError, match='Structured output is not supported by the model.'):
+    with pytest.raises(UserError, match='Native structured output is not supported by the model.'):
         await agent.run('What is the largest city in the user country?')
 
 

tests/test_logfire.py

@@ -289,6 +289,7 @@ async def my_ret(x: int) -> str:
                                 },
                                 'outer_typed_dict_key': None,
                                 'strict': None,
+                                'kind': 'function',
                             }
                         ],
                         'output_mode': 'text',

tests/test_streaming.py

@@ -613,18 +613,18 @@ def another_tool(y: int) -> int:
                         timestamp=IsNow(tz=timezone.utc),
                         tool_call_id=IsStr(),
                     ),
-                    RetryPromptPart(
-                        tool_name='unknown_tool',
-                        content="Unknown tool name: 'unknown_tool'. Available tools: regular_tool, another_tool, final_result",
-                        timestamp=IsNow(tz=timezone.utc),
-                        tool_call_id=IsStr(),
-                    ),
                     ToolReturnPart(
                         tool_name='regular_tool', content=42, timestamp=IsNow(tz=timezone.utc), tool_call_id=IsStr()
                     ),
                     ToolReturnPart(
                         tool_name='another_tool', content=2, timestamp=IsNow(tz=timezone.utc), tool_call_id=IsStr()
                     ),
+                    RetryPromptPart(
+                        content="Unknown tool name: 'unknown_tool'. Available tools: final_result, regular_tool, another_tool",
+                        tool_name='unknown_tool',
+                        tool_call_id=IsStr(),
+                        timestamp=IsNow(tz=timezone.utc),
+                    ),
                 ]
             ),
         ]
@@ -712,15 +712,15 @@ def another_tool(y: int) -> int:  # pragma: no cover
             ModelRequest(
                 parts=[
                     ToolReturnPart(
-                        tool_name='regular_tool',
-                        content='Tool not executed - a final result was already processed.',
+                        tool_name='final_result',
+                        content='Final result processed.',
                         tool_call_id=IsStr(),
                         timestamp=IsNow(tz=datetime.timezone.utc),
                         part_kind='tool-return',
                     ),
                     ToolReturnPart(
-                        tool_name='final_result',
-                        content='Final result processed.',
+                        tool_name='regular_tool',
+                        content='Tool not executed - a final result was already processed.',
                         tool_call_id=IsStr(),
                         timestamp=IsNow(tz=datetime.timezone.utc),
                         part_kind='tool-return',
@@ -733,10 +733,7 @@ def another_tool(y: int) -> int:  # pragma: no cover
                         part_kind='tool-return',
                     ),
                     RetryPromptPart(
-                        content='Unknown tool name: '
-                        "'unknown_tool'. Available tools: "
-                        'regular_tool, another_tool, '
-                        'final_result',
+                        content="Unknown tool name: 'unknown_tool'. Available tools: final_result, regular_tool, another_tool",
                         tool_name='unknown_tool',
                         tool_call_id=IsStr(),
                         timestamp=IsNow(tz=datetime.timezone.utc),
@@ -975,6 +972,13 @@ def known_tool(x: int) -> int:
 
     assert event_parts == snapshot(
         [
+            FunctionToolCallEvent(
+                part=ToolCallPart(
+                    tool_name='known_tool',
+                    args={'x': 5},
+                    tool_call_id=IsStr(),
+                )
+            ),
             FunctionToolCallEvent(
                 part=ToolCallPart(
                     tool_name='unknown_tool',
@@ -991,9 +995,6 @@ def known_tool(x: int) -> int:
                 ),
                 tool_call_id=IsStr(),
             ),
-            FunctionToolCallEvent(
-                part=ToolCallPart(tool_name='known_tool', args={'x': 5}, tool_call_id=IsStr()),
-            ),
             FunctionToolResultEvent(
                 result=ToolReturnPart(
                     tool_name='known_tool',
@@ -1003,13 +1004,6 @@ def known_tool(x: int) -> int:
                 ),
                 tool_call_id=IsStr(),
             ),
-            FunctionToolCallEvent(
-                part=ToolCallPart(
-                    tool_name='unknown_tool',
-                    args={'arg': 'value'},
-                    tool_call_id=IsStr(),
-                ),
-            ),
         ]
     )
 
@@ -1029,15 +1023,15 @@ def call_final_result_with_bad_data(messages: list[ModelMessage], info: AgentInf
 
     agent = Agent(FunctionModel(call_final_result_with_bad_data), output_type=OutputType)
 
-    event_parts: list[Any] = []
+    events: list[Any] = []
     async with agent.iter('test') as agent_run:
         async for node in agent_run:
             if Agent.is_call_tools_node(node):
                 async with node.stream(agent_run.ctx) as event_stream:
                     async for event in event_stream:
-                        event_parts.append(event)
+                        events.append(event)
 
-    assert event_parts == snapshot(
+    assert events == snapshot(
         [
             FunctionToolCallEvent(
                 part=ToolCallPart(
@@ -1047,9 +1041,16 @@ def call_final_result_with_bad_data(messages: list[ModelMessage], info: AgentInf
                 ),
             ),
             FunctionToolResultEvent(
-                result=ToolReturnPart(
+                result=RetryPromptPart(
+                    content=[
+                        {
+                            'type': 'missing',
+                            'loc': ('value',),
+                            'msg': 'Field required',
+                            'input': {'bad_value': 'invalid'},
+                        }
+                    ],
                     tool_name='final_result',
-                    content='Output tool not used - result failed validation.',
                     tool_call_id=IsStr(),
                     timestamp=IsNow(tz=timezone.utc),
                 ),