Add with agent.sequential_tool_calls(): contextmanager and use it in DBOSAgent (#2856)

Author: DouweM

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass, field, replace
 from typing import Any, Generic
 
@@ -16,6 +19,8 @@
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 from .usage import UsageLimits
 
+_sequential_tool_calls_ctx_var: ContextVar[bool] = ContextVar('sequential_tool_calls', default=False)
+
 
 @dataclass
 class ToolManager(Generic[AgentDepsT]):
@@ -30,6 +35,16 @@ class ToolManager(Generic[AgentDepsT]):
     failed_tools: set[str] = field(default_factory=set)
     """Names of tools that failed in this run step."""
 
+    @classmethod
+    @contextmanager
+    def sequential_tool_calls(cls) -> Iterator[None]:
+        """Run tool calls sequentially during the context."""
+        token = _sequential_tool_calls_ctx_var.set(True)
+        try:
+            yield
+        finally:
+            _sequential_tool_calls_ctx_var.reset(token)
+
     async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
         """Build a new tool manager for the next run step, carrying over the retries from the current run step."""
         if self.ctx is not None:
@@ -58,7 +73,9 @@ def tool_defs(self) -> list[ToolDefinition]:
 
     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)))
+        return _sequential_tool_calls_ctx_var.get() or 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."""

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -21,6 +21,7 @@
     result,
     usage as _usage,
 )
+from .._tool_manager import ToolManager
 from ..output import OutputDataT, OutputSpec
 from ..result import AgentStream, FinalResult, StreamedRunResult
 from ..run import AgentRun, AgentRunResult
@@ -714,6 +715,13 @@ def _infer_name(self, function_frame: FrameType | None) -> None:
                             self.name = name
                             return
 
+    @staticmethod
+    @contextmanager
+    def sequential_tool_calls() -> Iterator[None]:
+        """Run tool calls sequentially during the context."""
+        with ToolManager.sequential_tool_calls():
+            yield
+
     @staticmethod
     def is_model_request_node(
         node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],

pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_agent.py

@@ -13,7 +13,6 @@
     models,
     usage as _usage,
 )
-from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.agent import AbstractAgent, AgentRun, AgentRunResult, EventStreamHandler, RunOutputDataT, WrapperAgent
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.mcp import MCPServer
@@ -22,6 +21,7 @@
 from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import (
+    AgentDepsT,
     DeferredToolResults,
     RunContext,
     Tool,
@@ -218,7 +218,10 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
     @contextmanager
     def _dbos_overrides(self) -> Iterator[None]:
         # Override with DBOSModel and DBOSMCPServer in the toolsets.
-        with super().override(model=self._model, toolsets=self._toolsets, tools=[]):
+        with (
+            super().override(model=self._model, toolsets=self._toolsets, tools=[]),
+            self.sequential_tool_calls(),
+        ):
             yield
 
     @overload

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py

@@ -21,14 +21,14 @@
     models,
     usage as _usage,
 )
-from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.agent import AbstractAgent, AgentRun, AgentRunResult, EventStreamHandler, RunOutputDataT, WrapperAgent
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import Model
 from pydantic_ai.output import OutputDataT, OutputSpec
 from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import (
+    AgentDepsT,
     DeferredToolResults,
     RunContext,
     Tool,

tests/test_agent.py

@@ -5,7 +5,7 @@
 from collections.abc import AsyncIterable, Callable
 from dataclasses import dataclass, replace
 from datetime import timezone
-from typing import Any, Union
+from typing import Any, Literal, Union
 
 import httpx
 import pytest
@@ -4327,7 +4327,8 @@ async def call_tools_parallel(messages: list[ModelMessage], info: AgentInfo) ->
     assert result.output == snapshot('finished')
 
 
-def test_sequential_calls():
+@pytest.mark.parametrize('mode', ['argument', 'contextmanager'])
+def test_sequential_calls(mode: Literal['argument', 'contextmanager']):
     """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:
@@ -4355,31 +4356,37 @@ async def call_tools_sequential(messages: list[ModelMessage], info: AgentInfo) -
 
     integer_holder: int = 1
 
-    @sequential_toolset.tool(sequential=False)
+    @sequential_toolset.tool
     def call_first():
         nonlocal integer_holder
         assert integer_holder == 1
 
-    @sequential_toolset.tool(sequential=True)
+    @sequential_toolset.tool(sequential=mode == 'argument')
     def increment_integer_holder():
         nonlocal integer_holder
         integer_holder = 2
 
-    @sequential_toolset.tool()
+    @sequential_toolset.tool
     def requires_approval():
         from pydantic_ai.exceptions import ApprovalRequired
 
         raise ApprovalRequired()
 
-    @sequential_toolset.tool(sequential=False)
+    @sequential_toolset.tool
     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()
+
+    if mode == 'contextmanager':
+        with agent.sequential_tool_calls():
+            result = agent.run_sync()
+    else:
+        result = agent.run_sync()
+
     assert result.output == snapshot(
         DeferredToolRequests(approvals=[ToolCallPart(tool_name='requires_approval', tool_call_id=IsStr())])
     )

tests/test_dbos.py

@@ -132,10 +132,11 @@ def dbos() -> Generator[DBOS, Any, None]:
 def cleanup_test_sqlite_file() -> Iterator[None]:
     if os.path.exists(DBOS_SQLITE_FILE):
         os.remove(DBOS_SQLITE_FILE)  # pragma: lax no cover
-    yield
-
-    if os.path.exists(DBOS_SQLITE_FILE):
-        os.remove(DBOS_SQLITE_FILE)  # pragma: lax no cover
+    try:
+        yield
+    finally:
+        if os.path.exists(DBOS_SQLITE_FILE):
+            os.remove(DBOS_SQLITE_FILE)  # pragma: lax no cover
 
 
 model = OpenAIChatModel(
@@ -256,6 +257,7 @@ async def test_complex_agent_run_in_workflow(allow_model_requests: None, dbos: D
             'complex_agent__model.request_stream',
             'event_stream_handler',
             'event_stream_handler',
+            'event_stream_handler',
             'complex_agent__mcp_server__mcp.call_tool',
             'event_stream_handler',
             'complex_agent__mcp_server__mcp.get_tools',