Support sequential tool calling (#2718)

Co-authored-by: Douwe Maan <[email protected]>

Author: strawgate

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -2,6 +2,8 @@
 
 import asyncio
 import dataclasses
+import inspect
+from asyncio import Task
 from collections import defaultdict, deque
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
@@ -740,7 +742,6 @@ async def process_function_tools(  # noqa: C901
     deferred_tool_results: dict[str, DeferredToolResult] = {}
     if build_run_context(ctx).tool_call_approved and ctx.deps.tool_call_results is not None:
         deferred_tool_results = ctx.deps.tool_call_results
-
         # Deferred tool calls are "run" as well, by reading their value from the tool call results
         calls_to_run.extend(tool_calls_by_kind['external'])
         calls_to_run.extend(tool_calls_by_kind['unapproved'])
@@ -819,47 +820,65 @@ async def _call_tools(
     for call in tool_calls:
         yield _messages.FunctionToolCallEvent(call)
 
-    # Run all tool tasks in parallel
     with tracer.start_as_current_span(
         'running tools',
         attributes={
             'tools': [call.tool_name for call in tool_calls],
             'logfire.msg': f'running {len(tool_calls)} tool{"" if len(tool_calls) == 1 else "s"}',
         },
     ):
-        tasks = [
-            asyncio.create_task(
-                _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
-                name=call.tool_name,
-            )
-            for call in tool_calls
-        ]
-
-        pending = tasks
-        while pending:
-            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-            for task in done:
-                index = tasks.index(task)
-                try:
-                    tool_part, tool_user_part = task.result()
-                except exceptions.CallDeferred:
-                    deferred_calls_by_index[index] = 'external'
-                except exceptions.ApprovalRequired:
-                    deferred_calls_by_index[index] = 'unapproved'
-                else:
-                    yield _messages.FunctionToolResultEvent(tool_part)
 
-                    tool_parts_by_index[index] = tool_part
-                    if tool_user_part:
-                        user_parts_by_index[index] = tool_user_part
+        async def handle_call_or_result(
+            coro_or_task: Awaitable[
+                tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]
+            ]
+            | Task[tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]],
+            index: int,
+        ) -> _messages.HandleResponseEvent | None:
+            try:
+                tool_part, tool_user_part = (
+                    (await coro_or_task) if inspect.isawaitable(coro_or_task) else coro_or_task.result()
+                )
+            except exceptions.CallDeferred:
+                deferred_calls_by_index[index] = 'external'
+            except exceptions.ApprovalRequired:
+                deferred_calls_by_index[index] = 'unapproved'
+            else:
+                tool_parts_by_index[index] = tool_part
+                if tool_user_part:
+                    user_parts_by_index[index] = tool_user_part
+
+                return _messages.FunctionToolResultEvent(tool_part)
+
+        if tool_manager.should_call_sequentially(tool_calls):
+            for index, call in enumerate(tool_calls):
+                if event := await handle_call_or_result(
+                    _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
+                    index,
+                ):
+                    yield event
+
+        else:
+            tasks = [
+                asyncio.create_task(
+                    _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
+                    name=call.tool_name,
+                )
+                for call in tool_calls
+            ]
+
+            pending = tasks
+            while pending:
+                done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+                for task in done:
+                    index = tasks.index(task)
+                    if event := await handle_call_or_result(coro_or_task=task, index=index):
+                        yield event
 
     # 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(tool_parts_by_index):
-        output_parts.append(tool_parts_by_index[k])
-
-    for k in sorted(user_parts_by_index):
-        output_parts.append(user_parts_by_index[k])
+    output_parts.extend([tool_parts_by_index[k] for k in sorted(tool_parts_by_index)])
+    output_parts.extend([user_parts_by_index[k] for k in sorted(user_parts_by_index)])
 
     for k in sorted(deferred_calls_by_index):
         output_deferred_calls[deferred_calls_by_index[k]].append(tool_calls[k])

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -56,6 +56,10 @@ def tool_defs(self) -> list[ToolDefinition]:
 
         return [tool.tool_def for tool in self.tools.values()]
 
+    def should_call_sequentially(self, calls: list[ToolCallPart]) -> bool:
+        """Whether to require sequential tool calls for a list of tool calls."""
+        return any(tool_def.sequential for call in calls if (tool_def := self.get_tool_def(call.tool_name)))
+
     def get_tool_def(self, name: str) -> ToolDefinition | None:
         """Get the tool definition for a given tool name, or `None` if the tool is unknown."""
         if self.tools is None:

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -1119,6 +1119,7 @@ def tool_plain(
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
@@ -1164,6 +1165,7 @@ async def spam(ctx: RunContext[str]) -> float:
             schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
@@ -1180,6 +1182,7 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams
                 require_parameter_descriptions,
                 schema_generator,
                 strict,
+                sequential,
                 requires_approval,
             )
             return func_

pydantic_ai_slim/pydantic_ai/tools.py

@@ -253,6 +253,7 @@ class Tool(Generic[AgentDepsT]):
     docstring_format: DocstringFormat
     require_parameter_descriptions: bool
     strict: bool | None
+    sequential: bool
     requires_approval: bool
     function_schema: _function_schema.FunctionSchema
     """
@@ -274,6 +275,7 @@ def __init__(
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
         function_schema: _function_schema.FunctionSchema | None = None,
     ):
@@ -327,6 +329,7 @@ async def prep_my_tool(
             schema_generator: The JSON schema generator class to use. Defaults to `GenerateToolJsonSchema`.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
             function_schema: The function schema to use for the tool. If not provided, it will be generated.
@@ -347,6 +350,7 @@ async def prep_my_tool(
         self.docstring_format = docstring_format
         self.require_parameter_descriptions = require_parameter_descriptions
         self.strict = strict
+        self.sequential = sequential
         self.requires_approval = requires_approval
 
     @classmethod
@@ -357,6 +361,7 @@ def from_schema(
         description: str | None,
         json_schema: JsonSchemaValue,
         takes_ctx: bool = False,
+        sequential: bool = False,
     ) -> Self:
         """Creates a Pydantic tool from a function and a JSON schema.
 
@@ -370,6 +375,7 @@ def from_schema(
             json_schema: The schema for the function arguments
             takes_ctx: An optional boolean parameter indicating whether the function
                 accepts the context object as an argument.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
 
         Returns:
             A Pydantic tool that calls the function
@@ -389,6 +395,7 @@ def from_schema(
             name=name,
             description=description,
             function_schema=function_schema,
+            sequential=sequential,
         )
 
     @property
@@ -398,6 +405,7 @@ def tool_def(self):
             description=self.description,
             parameters_json_schema=self.function_schema.json_schema,
             strict=self.strict,
+            sequential=self.sequential,
         )
 
     async def prepare_tool_def(self, ctx: RunContext[AgentDepsT]) -> ToolDefinition | None:
@@ -466,6 +474,9 @@ class ToolDefinition:
     Note: this is currently only supported by OpenAI models.
     """
 
+    sequential: bool = False
+    """Whether this tool requires a sequential/serial execution environment."""
+
     kind: ToolKind = field(default='function')
     """The kind of tool:
 

pydantic_ai_slim/pydantic_ai/toolsets/function.py

@@ -97,6 +97,7 @@ def tool(
         require_parameter_descriptions: bool | None = None,
         schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> Callable[[ToolFuncEither[AgentDepsT, ToolParams]], ToolFuncEither[AgentDepsT, ToolParams]]: ...
 
@@ -112,6 +113,7 @@ def tool(
         require_parameter_descriptions: bool | None = None,
         schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
@@ -161,6 +163,7 @@ async def spam(ctx: RunContext[str], y: float) -> float:
                 If `None`, the default value is determined by the toolset.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
@@ -179,6 +182,7 @@ def tool_decorator(
                 require_parameter_descriptions,
                 schema_generator,
                 strict,
+                sequential,
                 requires_approval,
             )
             return func_
@@ -196,6 +200,7 @@ def add_function(
         require_parameter_descriptions: bool | None = None,
         schema_generator: type[GenerateJsonSchema] | None = None,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> None:
         """Add a function as a tool to the toolset.
@@ -222,6 +227,7 @@ def add_function(
                 If `None`, the default value is determined by the toolset.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
@@ -242,6 +248,7 @@ def add_function(
             require_parameter_descriptions=require_parameter_descriptions,
             schema_generator=schema_generator,
             strict=strict,
+            sequential=sequential,
             requires_approval=requires_approval,
         )
         self.add_tool(tool)

tests/test_agent.py

@@ -4293,6 +4293,65 @@ async def call_tools_parallel(messages: list[ModelMessage], info: AgentInfo) ->
     assert result.output == snapshot('finished')
 
 
+def test_sequential_calls():
+    """Test that tool calls are executed correctly when a `sequential` tool is present in the call."""
+
+    async def call_tools_sequential(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        return ModelResponse(
+            parts=[
+                ToolCallPart(tool_name='call_first'),
+                ToolCallPart(tool_name='call_first'),
+                ToolCallPart(tool_name='call_first'),
+                ToolCallPart(tool_name='call_first'),
+                ToolCallPart(tool_name='call_first'),
+                ToolCallPart(tool_name='call_first'),
+                ToolCallPart(tool_name='increment_integer_holder'),
+                ToolCallPart(tool_name='requires_approval'),
+                ToolCallPart(tool_name='call_second'),
+                ToolCallPart(tool_name='call_second'),
+                ToolCallPart(tool_name='call_second'),
+                ToolCallPart(tool_name='call_second'),
+                ToolCallPart(tool_name='call_second'),
+                ToolCallPart(tool_name='call_second'),
+                ToolCallPart(tool_name='call_second'),
+            ]
+        )
+
+    sequential_toolset = FunctionToolset()
+
+    integer_holder: int = 1
+
+    @sequential_toolset.tool(sequential=False)
+    def call_first():
+        nonlocal integer_holder
+        assert integer_holder == 1
+
+    @sequential_toolset.tool(sequential=True)
+    def increment_integer_holder():
+        nonlocal integer_holder
+        integer_holder = 2
+
+    @sequential_toolset.tool()
+    def requires_approval():
+        from pydantic_ai.exceptions import ApprovalRequired
+
+        raise ApprovalRequired()
+
+    @sequential_toolset.tool(sequential=False)
+    def call_second():
+        nonlocal integer_holder
+        assert integer_holder == 2
+
+    agent = Agent(
+        FunctionModel(call_tools_sequential), toolsets=[sequential_toolset], output_type=[str, DeferredToolRequests]
+    )
+    result = agent.run_sync()
+    assert result.output == snapshot(
+        DeferredToolRequests(approvals=[ToolCallPart(tool_name='requires_approval', tool_call_id=IsStr())])
+    )
+    assert integer_holder == 2
+
+
 def test_set_mcp_sampling_model():
     try:
         from pydantic_ai.mcp import MCPServerStdio

tests/test_logfire.py

@@ -385,6 +385,7 @@ async def my_ret(x: int) -> str:
                                 },
                                 'outer_typed_dict_key': None,
                                 'strict': None,
+                                'sequential': False,
                                 'kind': 'function',
                             }
                         ],
@@ -777,6 +778,7 @@ class MyOutput:
                                 'description': 'The final response which ends this conversation',
                                 'outer_typed_dict_key': None,
                                 'strict': None,
+                                'sequential': False,
                                 'kind': 'output',
                             }
                         ],