Allow dynamic in Temporal workflows

Author: qyinm

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py

@@ -14,6 +14,7 @@
 
 from ...exceptions import UserError
 from ._agent import TemporalAgent
+from ._function_toolset import TemporalFunctionToolset
 from ._logfire import LogfirePlugin
 from ._run_context import TemporalRunContext
 from ._toolset import TemporalWrapperToolset
@@ -26,6 +27,7 @@
     'AgentPlugin',
     'TemporalRunContext',
     'TemporalWrapperToolset',
+    'TemporalFunctionToolset',
     'PydanticAIWorkflow',
 ]
 

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py

@@ -242,10 +242,21 @@ def temporal_activities(self) -> list[Callable[..., Any]]:
         return self._temporal_activities
 
     @contextmanager
-    def _temporal_overrides(self) -> Iterator[None]:
-        # We reset tools here as the temporalized function toolset is already in self._toolsets.
+    def _temporal_overrides(self, toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None) -> Iterator[None]:
+        in_workflow = workflow.in_workflow()
+
+        if toolsets:
+            if in_workflow and any(not isinstance(t, TemporalWrapperToolset) for t in toolsets):
+                raise UserError(
+                    'Toolsets provided at runtime inside a Temporal workflow must be wrapped in a `TemporalWrapperToolset`.'
+                )
+            overridden_toolsets = [*self._toolsets, *toolsets]
+        else:
+            overridden_toolsets = list(self._toolsets)
+
+        # We reset tools here as the temporalized function toolset is already in overridden_toolsets.
         with (
-            super().override(model=self._model, toolsets=self._toolsets, tools=[]),
+            super().override(model=self._model, toolsets=overridden_toolsets, tools=[]),
             _utils.disable_threads(),
         ):
             temporal_active_token = self._temporal_overrides_active.set(True)
@@ -359,7 +370,7 @@ async def main():
                 'Event stream handler cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.'
             )
 
-        with self._temporal_overrides():
+        with self._temporal_overrides(toolsets=toolsets):
             return await super().run(
                 user_prompt,
                 output_type=output_type,
@@ -372,7 +383,7 @@ async def main():
                 usage_limits=usage_limits,
                 usage=usage,
                 infer_name=infer_name,
-                toolsets=toolsets,
+                toolsets=None,  # Toolsets are set via _temporal_overrides
                 builtin_tools=builtin_tools,
                 event_stream_handler=event_stream_handler or self.event_stream_handler,
                 **_deprecated_kwargs,
@@ -895,9 +906,9 @@ async def main():
                 raise UserError(
                     'Model cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.'
                 )
-            if toolsets is not None:
+            if toolsets is not None and any(not isinstance(t, TemporalWrapperToolset) for t in toolsets):
                 raise UserError(
-                    'Toolsets cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.'
+                    'Toolsets cannot be set at agent run time inside a Temporal workflow, unless they are wrapped in a `TemporalWrapperToolset`.'
                 )
 
         async with super().iter(
@@ -947,9 +958,9 @@ def override(
                 raise UserError(
                     'Model cannot be contextually overridden inside a Temporal workflow, it must be set at agent creation time.'
                 )
-            if _utils.is_set(toolsets):
+            if _utils.is_set(toolsets) and any(not isinstance(t, TemporalWrapperToolset) for t in toolsets):
                 raise UserError(
-                    'Toolsets cannot be contextually overridden inside a Temporal workflow, they must be set at agent creation time.'
+                    'Toolsets cannot be contextually overridden inside a Temporal workflow, unless they are wrapped in a `TemporalWrapperToolset`.'
                 )
             if _utils.is_set(tools):
                 raise UserError(

tests/test_temporal.py

@@ -1512,7 +1512,7 @@ async def test_temporal_agent_run_in_workflow_with_toolsets(allow_model_requests
         with workflow_raises(
             UserError,
             snapshot(
-                'Toolsets cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.'
+                'Toolsets provided at runtime inside a Temporal workflow must be wrapped in a `TemporalWrapperToolset`.'
             ),
         ):
             await client.execute_workflow(
@@ -1570,7 +1570,7 @@ async def test_temporal_agent_override_toolsets_in_workflow(allow_model_requests
         with workflow_raises(
             UserError,
             snapshot(
-                'Toolsets cannot be contextually overridden inside a Temporal workflow, they must be set at agent creation time.'
+                'Toolsets cannot be contextually overridden inside a Temporal workflow, unless they are wrapped in a `TemporalWrapperToolset`.'
             ),
         ):
             await client.execute_workflow(

tests/test_temporal_dynamic.py

@@ -0,0 +1,96 @@
+from datetime import timedelta
+
+import pytest
+from temporalio import workflow
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+from temporalio.workflow import ActivityConfig
+
+from pydantic_ai import Agent, FunctionToolset
+from pydantic_ai.durable_exec.temporal import PydanticAIPlugin, TemporalAgent, TemporalFunctionToolset
+from pydantic_ai.models.test import TestModel
+
+
+# 1. Define Tool
+def echo(x: str) -> str:
+    return f'echo: {x}'
+
+
+# 2. Create Toolset with specific ID
+# Using an explicit ID allows us to reference it later if needed,
+# though here we pass the toolset object directly.
+toolset = FunctionToolset(tools=[echo], id='my_tools')
+
+# 3. Wrap Toolset for Temporal (DouweM pattern)
+wrapped_toolset = TemporalFunctionToolset(
+    toolset,
+    activity_name_prefix='shared_tools',
+    activity_config=ActivityConfig(start_to_close_timeout=timedelta(minutes=1)),
+    tool_activity_config={},
+    deps_type=type(None),
+)
+
+# 4. Create base agent for model activity registration
+# This agent's activities will be registered in the Worker
+# We use a known name "test_agent" so the dynamic agent can share it.
+base_model = TestModel()
+base_agent = Agent(base_model, name='test_agent')
+base_temporal_agent = TemporalAgent(base_agent)
+
+
+# 5. Define Workflow
+@workflow.defn
+class DynamicToolWorkflow:
+    @workflow.run
+    async def run(self, user_prompt: str) -> str:
+        # Create agent dynamically within the workflow
+        # Note: We are using TestModel which mocks LLM behavior.
+        # We explicitly tell TestModel to call the 'echo' tool.
+        model = TestModel(call_tools=['echo'])
+
+        # We reuse the name "test_agent" so that the model activities
+        # (which are registered under that name) can be found.
+        # TEST: Revert to run-time passing
+        agent = Agent(
+            model,
+            name='test_agent',
+        )
+
+        temporal_agent = TemporalAgent(agent)
+
+        # Pass wrapped toolset at runtime
+        result = await temporal_agent.run(user_prompt, toolsets=[wrapped_toolset])
+        return result.output
+
+
+# 6. Test
+pytestmark = pytest.mark.anyio
+
+
+async def test_dynamic_tool_registration():
+    """Test passing a `TemporalWrapperToolset` at runtime to a `TemporalAgent` within a workflow."""
+    env = await WorkflowEnvironment.start_local()  # type: ignore[reportUnknownMemberType]
+    async with env:
+        async with Worker(
+            env.client,
+            task_queue='test-queue',
+            workflows=[DynamicToolWorkflow],
+            # Register activities from both base agent and shared toolset
+            activities=[
+                *base_temporal_agent.temporal_activities,
+                *wrapped_toolset.temporal_activities,
+            ],
+            plugins=[PydanticAIPlugin()],
+        ):
+            result = await env.client.execute_workflow(
+                DynamicToolWorkflow.run,
+                args=['test prompt'],
+                id='test-workflow-run',
+                task_queue='test-queue',
+            )
+
+            # Verify tool was called successfully
+            # TestModel generates random args, so we just verify echo was called
+            # "echo" is the tool return value format: "echo: {arg}"
+            assert 'echo' in result
+            assert 'echo:' in result