From b37aa26b2afa2469d0f65aa9b7226787588e61bc Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Tue, 21 Oct 2025 23:22:14 +0000 Subject: [PATCH 1/4] Bump temporalio and use SimplePlugin --- .../durable_exec/temporal/__init__.py | 161 ++++++------------ .../durable_exec/temporal/_logfire.py | 28 ++- pyproject.toml | 1 + uv.lock | 14 +- 4 files changed, 72 insertions(+), 132 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py index b982527345..b711157c89 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py @@ -1,24 +1,14 @@ from __future__ import annotations import warnings -from collections.abc import AsyncIterator, Callable, Sequence -from contextlib import AbstractAsyncContextManager from dataclasses import replace from typing import Any from pydantic.errors import PydanticUserError -from temporalio.client import ClientConfig, Plugin as ClientPlugin, WorkflowHistory from temporalio.contrib.pydantic import PydanticPayloadConverter, pydantic_data_converter from temporalio.converter import DataConverter, DefaultPayloadConverter -from temporalio.service import ConnectConfig, ServiceClient -from temporalio.worker import ( - Plugin as WorkerPlugin, - Replayer, - ReplayerConfig, - Worker, - WorkerConfig, - WorkflowReplayResult, -) +from temporalio.plugin import SimplePlugin +from temporalio.worker import WorkflowRunner from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner from ...exceptions import UserError @@ -37,102 +27,61 @@ ] -class PydanticAIPlugin(ClientPlugin, WorkerPlugin): +def _data_converter(converter: DataConverter | None) -> DataConverter: + if converter and converter.payload_converter_class not in ( + DefaultPayloadConverter, + PydanticPayloadConverter, + ): + warnings.warn( # pragma: no cover + 'A non-default Temporal data converter was used which has been replaced with the Pydantic data converter.' + ) + + return pydantic_data_converter + + +def _workflow_runner(runner: WorkflowRunner | None) -> WorkflowRunner: + if not runner: + raise ValueError('No WorkflowRunner provided to the Pydantic AI plugin.') + + if isinstance(runner, SandboxedWorkflowRunner): + return replace( + runner, + restrictions=runner.restrictions.with_passthrough_modules( + 'pydantic_ai', + 'pydantic', + 'pydantic_core', + 'logfire', + 'rich', + 'httpx', + 'anyio', + 'httpcore', + # Imported inside `logfire._internal.json_encoder` when running `logfire.info` inside an activity with attributes to serialize + 'attrs', + # Imported inside `logfire._internal.json_schema` when running `logfire.info` inside an activity with attributes to serialize + 'numpy', + 'pandas', + ), + ) + return runner + + +class PydanticAIPlugin(SimplePlugin): """Temporal client and worker plugin for Pydantic AI.""" - def init_client_plugin(self, next: ClientPlugin) -> None: - self.next_client_plugin = next - - def init_worker_plugin(self, next: WorkerPlugin) -> None: - self.next_worker_plugin = next - - def configure_client(self, config: ClientConfig) -> ClientConfig: - config['data_converter'] = self._get_new_data_converter(config.get('data_converter')) - return self.next_client_plugin.configure_client(config) - - def configure_worker(self, config: WorkerConfig) -> WorkerConfig: - runner = config.get('workflow_runner') # pyright: ignore[reportUnknownMemberType] - if isinstance(runner, SandboxedWorkflowRunner): # pragma: no branch - config['workflow_runner'] = replace( - runner, - restrictions=runner.restrictions.with_passthrough_modules( - 'pydantic_ai', - 'pydantic', - 'pydantic_core', - 'logfire', - 'rich', - 'httpx', - 'anyio', - 'httpcore', - # Imported inside `logfire._internal.json_encoder` when running `logfire.info` inside an activity with attributes to serialize - 'attrs', - # Imported inside `logfire._internal.json_schema` when running `logfire.info` inside an activity with attributes to serialize - 'numpy', - 'pandas', - ), - ) - - config['workflow_failure_exception_types'] = [ - *config.get('workflow_failure_exception_types', []), # pyright: ignore[reportUnknownMemberType] - UserError, - PydanticUserError, - ] - - return self.next_worker_plugin.configure_worker(config) - - async def connect_service_client(self, config: ConnectConfig) -> ServiceClient: - return await self.next_client_plugin.connect_service_client(config) - - async def run_worker(self, worker: Worker) -> None: - await self.next_worker_plugin.run_worker(worker) - - def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig: # pragma: no cover - config['data_converter'] = self._get_new_data_converter(config.get('data_converter')) # pyright: ignore[reportUnknownMemberType] - return self.next_worker_plugin.configure_replayer(config) - - def run_replayer( - self, - replayer: Replayer, - histories: AsyncIterator[WorkflowHistory], - ) -> AbstractAsyncContextManager[AsyncIterator[WorkflowReplayResult]]: # pragma: no cover - return self.next_worker_plugin.run_replayer(replayer, histories) - - def _get_new_data_converter(self, converter: DataConverter | None) -> DataConverter: - if converter and converter.payload_converter_class not in ( - DefaultPayloadConverter, - PydanticPayloadConverter, - ): - warnings.warn( # pragma: no cover - 'A non-default Temporal data converter was used which has been replaced with the Pydantic data converter.' - ) - - return pydantic_data_converter - - -class AgentPlugin(WorkerPlugin): - """Temporal worker plugin for a specific Pydantic AI agent.""" - - def __init__(self, agent: TemporalAgent[Any, Any]): - self.agent = agent - - def init_worker_plugin(self, next: WorkerPlugin) -> None: - self.next_worker_plugin = next + def __init__(self): + super().__init__( # type: ignore[reportUnknownMemberType] + name='PydanticAIPlugin', + data_converter=_data_converter, + workflow_runner=_workflow_runner, + workflow_failure_exception_types=[UserError, PydanticUserError], + ) - def configure_worker(self, config: WorkerConfig) -> WorkerConfig: - activities: Sequence[Callable[..., Any]] = config.get('activities', []) # pyright: ignore[reportUnknownMemberType] - # Activities are checked for name conflicts by Temporal. - config['activities'] = [*activities, *self.agent.temporal_activities] - return self.next_worker_plugin.configure_worker(config) - async def run_worker(self, worker: Worker) -> None: - await self.next_worker_plugin.run_worker(worker) - - def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig: # pragma: no cover - return self.next_worker_plugin.configure_replayer(config) +class AgentPlugin(SimplePlugin): + """Temporal worker plugin for a specific Pydantic AI agent.""" - def run_replayer( - self, - replayer: Replayer, - histories: AsyncIterator[WorkflowHistory], - ) -> AbstractAsyncContextManager[AsyncIterator[WorkflowReplayResult]]: # pragma: no cover - return self.next_worker_plugin.run_replayer(replayer, histories) + def __init__(self, agent: TemporalAgent[Any, Any]): + super().__init__( # type: ignore[reportUnknownMemberType] + name='AgentPlugin', + activities=agent.temporal_activities, + ) diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_logfire.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_logfire.py index 055567e4d6..1c1c8f8a08 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_logfire.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_logfire.py @@ -1,9 +1,9 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING -from temporalio.client import ClientConfig, Plugin as ClientPlugin +from temporalio.plugin import SimplePlugin from temporalio.runtime import OpenTelemetryConfig, Runtime, TelemetryConfig from temporalio.service import ConnectConfig, ServiceClient @@ -19,12 +19,14 @@ def _default_setup_logfire() -> Logfire: return instance -class LogfirePlugin(ClientPlugin): +class LogfirePlugin(SimplePlugin): """Temporal client plugin for Logfire.""" def __init__(self, setup_logfire: Callable[[], Logfire] = _default_setup_logfire, *, metrics: bool = True): try: import logfire # noqa: F401 # pyright: ignore[reportUnusedImport] + from opentelemetry.trace import get_tracer + from temporalio.contrib.opentelemetry import TracingInterceptor except ImportError as _import_error: raise ImportError( 'Please install the `logfire` package to use the Logfire plugin, ' @@ -34,18 +36,14 @@ def __init__(self, setup_logfire: Callable[[], Logfire] = _default_setup_logfire self.setup_logfire = setup_logfire self.metrics = metrics - def init_client_plugin(self, next: ClientPlugin) -> None: - self.next_client_plugin = next + super().__init__( # type: ignore[reportUnknownMemberType] + name='LogfirePlugin', + client_interceptors=[TracingInterceptor(get_tracer('temporalio'))], + ) - def configure_client(self, config: ClientConfig) -> ClientConfig: - from opentelemetry.trace import get_tracer - from temporalio.contrib.opentelemetry import TracingInterceptor - - interceptors = config.get('interceptors', []) - config['interceptors'] = [*interceptors, TracingInterceptor(get_tracer('temporalio'))] - return self.next_client_plugin.configure_client(config) - - async def connect_service_client(self, config: ConnectConfig) -> ServiceClient: + async def connect_service_client( + self, config: ConnectConfig, next: Callable[[ConnectConfig], Awaitable[ServiceClient]] + ) -> ServiceClient: logfire = self.setup_logfire() if self.metrics: @@ -60,4 +58,4 @@ async def connect_service_client(self, config: ConnectConfig) -> ServiceClient: telemetry=TelemetryConfig(metrics=OpenTelemetryConfig(url=metrics_url, headers=headers)) ) - return await self.next_client_plugin.connect_service_client(config) + return await next(config) diff --git a/pyproject.toml b/pyproject.toml index b09a172045..fd292d1d6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,7 @@ pydantic-ai-slim = { workspace = true } pydantic-evals = { workspace = true } pydantic-graph = { workspace = true } pydantic-ai-examples = { workspace = true } +temporalio = { git = "https://github.com/temporalio/sdk-python.git", rev = "main" } [tool.uv.workspace] members = [ diff --git a/uv.lock b/uv.lock index 486ea9bbd3..0075661bde 100644 --- a/uv.lock +++ b/uv.lock @@ -3921,7 +3921,7 @@ requires-dist = [ { name = "rich", marker = "extra == 'cli'", specifier = ">=13" }, { name = "starlette", marker = "extra == 'ag-ui'", specifier = ">=0.45.3" }, { name = "tavily-python", marker = "extra == 'tavily'", specifier = ">=0.5.0" }, - { name = "temporalio", marker = "extra == 'temporal'", specifier = "==1.18.0" }, + { name = "temporalio", marker = "extra == 'temporal'", git = "https://github.com/temporalio/sdk-python.git?rev=main" }, { name = "tenacity", marker = "extra == 'retries'", specifier = ">=8.2.3" }, { name = "typing-inspection", specifier = ">=0.4.0" }, ] @@ -4938,8 +4938,8 @@ wheels = [ [[package]] name = "temporalio" -version = "1.18.0" -source = { registry = "https://pypi.org/simple" } +version = "1.18.1" +source = { git = "https://github.com/temporalio/sdk-python.git?rev=main#f03ddc2136c56acda77387627f9b379166d06d03" } dependencies = [ { name = "nexus-rpc" }, { name = "protobuf" }, @@ -4947,14 +4947,6 @@ dependencies = [ { name = "types-protobuf" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7e/20/b52c96b37bf00ead6e8a4a197075770ebad516db765cc3abca8396de0689/temporalio-1.18.0.tar.gz", hash = "sha256:7ff7f833eb1e7697084b4ed9d86c3167cbff1ec77f1b40df774313a5d0fd5f6d", size = 1781572, upload-time = "2025-09-19T23:40:52.511Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2f/28/c5a4ee259748450ac0765837f8c78cbfa36800264158d98bd2cde4496d87/temporalio-1.18.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:ac5d30d8b010c9b042065ea1259da7638db1a0a25e81ee4be0671a393ed329c5", size = 12734753, upload-time = "2025-09-19T23:40:06.575Z" }, - { url = "https://files.pythonhosted.org/packages/be/94/24bd903b5594420a4d131bfa3de965313f9a409af77b47e9a9a56d85bb9e/temporalio-1.18.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:19315d192247230c9bd7c60a566c2b3a80ad4d9de891c6aa13df63d72d3ec169", size = 12323141, upload-time = "2025-09-19T23:40:16.817Z" }, - { url = "https://files.pythonhosted.org/packages/6d/76/82415b43c68e2c6bb3a85e8800555d206767815088c8cad0ade9a06bd7ac/temporalio-1.18.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a023b25033e48b2e43f623a78737047a45b8cb553f69f457d09272fce5c723da", size = 12694061, upload-time = "2025-09-19T23:40:26.388Z" }, - { url = "https://files.pythonhosted.org/packages/41/60/176a3224c2739fee270052dd9224ae36370c4e13d2ab1bb96a2f9bbb513c/temporalio-1.18.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:695211dddbcffc20077d5b3b9a9b41bd09f60393c4ff211bcc7d6d895d607cc1", size = 12879404, upload-time = "2025-09-19T23:40:37.487Z" }, - { url = "https://files.pythonhosted.org/packages/e3/8d/e3809b356262d1d398d8cbb78df1e19d460c0a89e6ab64ca8d9c05d5fe5a/temporalio-1.18.0-cp39-abi3-win_amd64.whl", hash = "sha256:e3f691bd0a01a22c0fe40e87b6236cc8a292628e3a5a490880d1bf94709249c9", size = 13088041, upload-time = "2025-09-19T23:40:49.469Z" }, -] [[package]] name = "tenacity" From 87d666233aa015b6a87b28099839d00131aaf6ed Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 14 Nov 2025 16:53:07 +0000 Subject: [PATCH 2/4] Bump temporalio to v1.19.0 --- pydantic_ai_slim/pyproject.toml | 2 +- pyproject.toml | 1 - uv.lock | 13 ++++++++++--- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index a1f917295d..6b0658b882 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/pydantic_ai_slim/pyproject.toml @@ -106,7 +106,7 @@ ag-ui = ["ag-ui-protocol>=0.1.8", "starlette>=0.45.3"] # Retries retries = ["tenacity>=8.2.3"] # Temporal -temporal = ["temporalio==1.18.2"] +temporal = ["temporalio==1.19.0"] # DBOS dbos = ["dbos>=1.14.0"] # Prefect diff --git a/pyproject.toml b/pyproject.toml index 08acf482dd..3c13afdece 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,7 +75,6 @@ pydantic-ai-slim = { workspace = true } pydantic-evals = { workspace = true } pydantic-graph = { workspace = true } pydantic-ai-examples = { workspace = true } -temporalio = { git = "https://github.com/temporalio/sdk-python.git", rev = "main" } [tool.uv.workspace] members = [ diff --git a/uv.lock b/uv.lock index a5602b11c9..c02ca155ea 100644 --- a/uv.lock +++ b/uv.lock @@ -5660,7 +5660,7 @@ requires-dist = [ { name = "starlette", marker = "extra == 'ag-ui'", specifier = ">=0.45.3" }, { name = "starlette", marker = "extra == 'ui'", specifier = ">=0.45.3" }, { name = "tavily-python", marker = "extra == 'tavily'", specifier = ">=0.5.0" }, - { name = "temporalio", marker = "extra == 'temporal'", git = "https://github.com/temporalio/sdk-python.git?rev=main" }, + { name = "temporalio", marker = "extra == 'temporal'", specifier = "==1.19.0" }, { name = "tenacity", marker = "extra == 'retries'", specifier = ">=8.2.3" }, { name = "torch", marker = "(platform_machine != 'x86_64' and extra == 'outlines-transformers') or (sys_platform != 'darwin' and extra == 'outlines-transformers')" }, { name = "torch", marker = "(platform_machine != 'x86_64' and extra == 'outlines-vllm-offline') or (sys_platform != 'darwin' and extra == 'outlines-vllm-offline')" }, @@ -7323,8 +7323,8 @@ wheels = [ [[package]] name = "temporalio" -version = "1.18.1" -source = { git = "https://github.com/temporalio/sdk-python.git?rev=main#f03ddc2136c56acda77387627f9b379166d06d03" } +version = "1.19.0" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nexus-rpc" }, { name = "protobuf" }, @@ -7332,6 +7332,13 @@ dependencies = [ { name = "types-protobuf" }, { name = "typing-extensions" }, ] +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/92/0775d831fa245d61b74db2059d5a24a04cef0532ed2c48310a5ab007de9c/temporalio-1.19.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c2d6d5cad8aec56e048705aa4f0bab83fec15343757ea7acf8504f2e0c289b60", size = 13175255, upload-time = "2025-11-13T22:35:54.22Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e1/2a818fefc0023eb132bfff1a03440bcaff154d4d97445ef88a40c23c20c8/temporalio-1.19.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:d85c89018cba9471ce529d90c9cee5bcc31790fd64176b9ada32cc76440f8d73", size = 12854549, upload-time = "2025-11-13T22:35:57.217Z" }, + { url = "https://files.pythonhosted.org/packages/ff/78/fe5c8c9b112b38e01aba845335df17a8bbfd60a434ffe3c1c4737ced40a0/temporalio-1.19.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f772f0698d60f808bc3c4a055fb53e40d757fa646411845b911863eebbf0549d", size = 13237772, upload-time = "2025-11-13T22:36:00.511Z" }, + { url = "https://files.pythonhosted.org/packages/d9/82/be0fd31119651f518f8db8685fd61976d9d5bbecf3b562d51f13a6442a17/temporalio-1.19.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f706c8f49771daf342ac8daa8ed07f4124fae943177f9feef458a1255aee717c", size = 13374621, upload-time = "2025-11-13T22:36:03.431Z" }, + { url = "https://files.pythonhosted.org/packages/d8/94/18f6ae06ffd91507ded9111af1041146a5ba4b56e9256520c5ce82629fc4/temporalio-1.19.0-cp310-abi3-win_amd64.whl", hash = "sha256:162459c293553be39994f20c635a132f7332ae71bd7ba4042f8473701fcf1c7c", size = 14256891, upload-time = "2025-11-13T22:36:06.778Z" }, +] [[package]] name = "tenacity" From cf3688827a78ea204446e6b5726a1a9df11df255 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 14 Nov 2025 17:20:36 +0000 Subject: [PATCH 3/4] remove unnecessary pyright ignores --- .../durable_exec/temporal/_agent.py | 2 +- .../temporal/_function_toolset.py | 2 +- .../pydantic_ai/durable_exec/temporal/_mcp.py | 4 +- .../durable_exec/temporal/_model.py | 4 +- tests/test_temporal.py | 48 +++++++++---------- 5 files changed, 30 insertions(+), 30 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py index 82ec0e76fa..6e964c8d08 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py @@ -219,7 +219,7 @@ async def _call_event_stream_handler_activity( ) -> None: serialized_run_context = self.run_context_type.serialize_run_context(ctx) async for event in stream: - await workflow.execute_activity( # pyright: ignore[reportUnknownMemberType] + await workflow.execute_activity( activity=self.event_stream_handler_activity, args=[ _EventStreamHandlerParams( diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py index dd1f8c1ee3..05bc3f5f2c 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py @@ -81,7 +81,7 @@ async def call_tool( tool_activity_config = self.activity_config | tool_activity_config serialized_run_context = self.run_context_type.serialize_run_context(ctx) return self._unwrap_call_tool_result( - await workflow.execute_activity( # pyright: ignore[reportUnknownMemberType] + await workflow.execute_activity( activity=self.call_tool_activity, args=[ CallToolParams( diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp.py index ab7f52d63e..aecb220a28 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp.py @@ -108,7 +108,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ return await super().get_tools(ctx) serialized_run_context = self.run_context_type.serialize_run_context(ctx) - tool_defs = await workflow.execute_activity( # pyright: ignore[reportUnknownMemberType] + tool_defs = await workflow.execute_activity( activity=self.get_tools_activity, args=[ _GetToolsParams(serialized_run_context=serialized_run_context), @@ -131,7 +131,7 @@ async def call_tool( tool_activity_config = self.activity_config | self.tool_activity_config.get(name, {}) serialized_run_context = self.run_context_type.serialize_run_context(ctx) return self._unwrap_call_tool_result( - await workflow.execute_activity( # pyright: ignore[reportUnknownMemberType] + await workflow.execute_activity( activity=self.call_tool_activity, args=[ CallToolParams( diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py index 8a97253ee1..794f36ba16 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py @@ -130,7 +130,7 @@ async def request( self._validate_model_request_parameters(model_request_parameters) - return await workflow.execute_activity( # pyright: ignore[reportUnknownMemberType] + return await workflow.execute_activity( activity=self.request_activity, arg=_RequestParams( messages=messages, @@ -168,7 +168,7 @@ async def request_stream( self._validate_model_request_parameters(model_request_parameters) serialized_run_context = self.run_context_type.serialize_run_context(run_context) - response = await workflow.execute_activity( # pyright: ignore[reportUnknownMemberType] + response = await workflow.execute_activity( activity=self.request_stream_activity, args=[ _RequestParams( diff --git a/tests/test_temporal.py b/tests/test_temporal.py index e6f655be14..938e123077 100644 --- a/tests/test_temporal.py +++ b/tests/test_temporal.py @@ -212,7 +212,7 @@ async def test_simple_agent_run_in_workflow(allow_model_requests: None, client: workflows=[SimpleAgentWorkflow], plugins=[AgentPlugin(simple_temporal_agent)], ): - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( SimpleAgentWorkflow.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflow.__name__, @@ -320,7 +320,7 @@ async def test_complex_agent_run_in_workflow( workflows=[ComplexAgentWorkflow], plugins=[AgentPlugin(complex_temporal_agent)], ): - output = await client_with_logfire.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client_with_logfire.execute_workflow( ComplexAgentWorkflow.run, args=[ 'Tell me: the capital of the country; the weather there; the product name', @@ -1037,7 +1037,7 @@ async def test_multiple_agents(allow_model_requests: None, client: Client): workflows=[SimpleAgentWorkflow, ComplexAgentWorkflow], plugins=[AgentPlugin(simple_temporal_agent), AgentPlugin(complex_temporal_agent)], ): - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( SimpleAgentWorkflow.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflow.__name__, @@ -1045,7 +1045,7 @@ async def test_multiple_agents(allow_model_requests: None, client: Client): ) assert output == snapshot('The capital of Mexico is Mexico City.') - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( ComplexAgentWorkflow.run, args=[ 'Tell me: the capital of the country; the weather there; the product name', @@ -1239,7 +1239,7 @@ async def test_temporal_agent_run_sync_in_workflow(allow_model_requests: None, c UserError, snapshot('`agent.run_sync()` cannot be used inside a Temporal workflow. Use `await agent.run()` instead.'), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithRunSync.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithRunSync.__name__, @@ -1269,7 +1269,7 @@ async def test_temporal_agent_run_stream_in_workflow(allow_model_requests: None, '`agent.run_stream()` cannot be used inside a Temporal workflow. Set an `event_stream_handler` on the agent and use `agent.run()` instead.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithRunStream.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithRunStream.__name__, @@ -1297,7 +1297,7 @@ async def test_temporal_agent_run_stream_events_in_workflow(allow_model_requests '`agent.run_stream_events()` cannot be used inside a Temporal workflow. Set an `event_stream_handler` on the agent and use `agent.run()` instead.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithRunStreamEvents.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithRunStreamEvents.__name__, @@ -1328,7 +1328,7 @@ async def test_temporal_agent_iter_in_workflow(allow_model_requests: None, clien '`agent.iter()` cannot be used inside a Temporal workflow. Set an `event_stream_handler` on the agent and use `agent.run()` instead.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithIter.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithIter.__name__, @@ -1364,7 +1364,7 @@ async def test_temporal_agent_run_in_workflow_with_event_stream_handler(allow_mo 'Event stream handler cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithEventStreamHandler.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithEventStreamHandler.__name__, @@ -1393,7 +1393,7 @@ async def test_temporal_agent_run_in_workflow_with_model(allow_model_requests: N 'Model cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithRunModel.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithRunModel.__name__, @@ -1422,7 +1422,7 @@ async def test_temporal_agent_run_in_workflow_with_toolsets(allow_model_requests 'Toolsets cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithRunToolsets.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithRunToolsets.__name__, @@ -1451,7 +1451,7 @@ async def test_temporal_agent_override_model_in_workflow(allow_model_requests: N 'Model cannot be contextually overridden inside a Temporal workflow, it must be set at agent creation time.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithOverrideModel.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithOverrideModel.__name__, @@ -1480,7 +1480,7 @@ async def test_temporal_agent_override_toolsets_in_workflow(allow_model_requests 'Toolsets cannot be contextually overridden inside a Temporal workflow, they must be set at agent creation time.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithOverrideToolsets.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithOverrideToolsets.__name__, @@ -1509,7 +1509,7 @@ async def test_temporal_agent_override_tools_in_workflow(allow_model_requests: N 'Tools cannot be contextually overridden inside a Temporal workflow, they must be set at agent creation time.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( SimpleAgentWorkflowWithOverrideTools.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithOverrideTools.__name__, @@ -1533,7 +1533,7 @@ async def test_temporal_agent_override_deps_in_workflow(allow_model_requests: No workflows=[SimpleAgentWorkflowWithOverrideDeps], plugins=[AgentPlugin(simple_temporal_agent)], ): - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( SimpleAgentWorkflowWithOverrideDeps.run, args=['What is the capital of Mexico?'], id=SimpleAgentWorkflowWithOverrideDeps.__name__, @@ -1577,7 +1577,7 @@ async def test_temporal_agent_sync_tool_activity_disabled(allow_model_requests: "Temporal activity config for tool 'get_weather' has been explicitly set to `False` (activity disabled), but non-async tools are run in threads which are not supported outside of an activity. Make the tool function async instead." ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( AgentWorkflowWithSyncToolActivityDisabled.run, args=['What is the weather in Mexico City?'], id=AgentWorkflowWithSyncToolActivityDisabled.__name__, @@ -1627,7 +1627,7 @@ async def test_temporal_model_stream_direct(client: Client): 'A Temporal model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. Set an `event_stream_handler` on the agent and use `agent.run()` instead.' ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( DirectStreamWorkflow.run, args=['What is the capital of Mexico?'], id=DirectStreamWorkflow.__name__, @@ -1668,7 +1668,7 @@ async def test_temporal_agent_with_unserializable_deps_type(allow_model_requests "The `deps` object failed to be serialized. Temporal requires all objects that are passed to activities to be serializable using Pydantic's `TypeAdapter`." ), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( UnserializableDepsAgentWorkflow.run, args=['What is the model name?'], id=UnserializableDepsAgentWorkflow.__name__, @@ -1787,7 +1787,7 @@ async def test_temporal_agent_with_hitl_tool(allow_model_requests: None, client: workflows=[HitlAgentWorkflow], plugins=[AgentPlugin(hitl_temporal_agent)], ): - workflow = await client.start_workflow( # pyright: ignore[reportUnknownMemberType] + workflow = await client.start_workflow( HitlAgentWorkflow.run, args=['Delete the file `.env` and create `test.txt`'], id=HitlAgentWorkflow.__name__, @@ -1933,7 +1933,7 @@ async def test_temporal_agent_with_model_retry(allow_model_requests: None, clien workflows=[ModelRetryWorkflow], plugins=[AgentPlugin(model_retry_temporal_agent)], ): - workflow = await client.start_workflow( # pyright: ignore[reportUnknownMemberType] + workflow = await client.start_workflow( ModelRetryWorkflow.run, args=['What is the weather in CDMX?'], id=ModelRetryWorkflow.__name__, @@ -2082,7 +2082,7 @@ async def test_custom_model_settings(allow_model_requests: None, client: Client) workflows=[SettingsAgentWorkflow], plugins=[AgentPlugin(settings_temporal_agent)], ): - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( SettingsAgentWorkflow.run, args=['Give me those settings'], id=SettingsAgentWorkflow.__name__, @@ -2116,7 +2116,7 @@ async def test_image_agent(allow_model_requests: None, client: Client): UserError, snapshot('Image output is not supported with Temporal because of the 2MB payload size limit.'), ): - await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + await client.execute_workflow( ImageAgentWorkflow.run, args=['Generate an image of an axolotl.'], id=ImageAgentWorkflow.__name__, @@ -2165,7 +2165,7 @@ async def test_web_search_agent_run_in_workflow(allow_model_requests: None, clie workflows=[WebSearchAgentWorkflow], plugins=[AgentPlugin(web_search_temporal_agent)], ): - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( WebSearchAgentWorkflow.run, args=['In one sentence, what is the top news story in my country today?'], id=WebSearchAgentWorkflow.__name__, @@ -2219,7 +2219,7 @@ async def test_fastmcp_toolset(allow_model_requests: None, client: Client): workflows=[FastMCPAgentWorkflow], plugins=[AgentPlugin(fastmcp_temporal_agent)], ): - output = await client.execute_workflow( # pyright: ignore[reportUnknownMemberType] + output = await client.execute_workflow( FastMCPAgentWorkflow.run, args=['Can you tell me more about the pydantic/pydantic-ai repo? Keep your answer short'], id=FastMCPAgentWorkflow.__name__, From cfe2b78d58dfeca27dfb741bb5061f2793b9da1a Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 14 Nov 2025 17:40:06 +0000 Subject: [PATCH 4/4] coverage --- .../durable_exec/temporal/__init__.py | 49 ++++++++++--------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py index e5e5d3246c..9611565fde 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/__init__.py @@ -52,30 +52,31 @@ def _data_converter(converter: DataConverter | None) -> DataConverter: def _workflow_runner(runner: WorkflowRunner | None) -> WorkflowRunner: if not runner: - raise ValueError('No WorkflowRunner provided to the Pydantic AI plugin.') - - if isinstance(runner, SandboxedWorkflowRunner): - return replace( - runner, - restrictions=runner.restrictions.with_passthrough_modules( - 'pydantic_ai', - 'pydantic', - 'pydantic_core', - 'logfire', - 'rich', - 'httpx', - 'anyio', - 'httpcore', - # Used by fastmcp via py-key-value-aio - 'beartype', - # Imported inside `logfire._internal.json_encoder` when running `logfire.info` inside an activity with attributes to serialize - 'attrs', - # Imported inside `logfire._internal.json_schema` when running `logfire.info` inside an activity with attributes to serialize - 'numpy', - 'pandas', - ), - ) - return runner + raise ValueError('No WorkflowRunner provided to the Pydantic AI plugin.') # pragma: no cover + + if not isinstance(runner, SandboxedWorkflowRunner): + return runner # pragma: no cover + + return replace( + runner, + restrictions=runner.restrictions.with_passthrough_modules( + 'pydantic_ai', + 'pydantic', + 'pydantic_core', + 'logfire', + 'rich', + 'httpx', + 'anyio', + 'httpcore', + # Used by fastmcp via py-key-value-aio + 'beartype', + # Imported inside `logfire._internal.json_encoder` when running `logfire.info` inside an activity with attributes to serialize + 'attrs', + # Imported inside `logfire._internal.json_schema` when running `logfire.info` inside an activity with attributes to serialize + 'numpy', + 'pandas', + ), + ) class PydanticAIPlugin(SimplePlugin):