Address review comments

Author: desertaxle

docs/durable_execution/prefect.md

@@ -61,7 +61,7 @@ Any agent can be wrapped in a [`PrefectAgent`][pydantic_ai.durable_exec.prefect.
 * Wraps [tool calls](../tools.md) as Prefect tasks (configurable per-tool).
 * Wraps [MCP communication](../mcp/client.md) as Prefect tasks.
 
-Event stream handlers are **not automatically wrapped** by Prefect. If they involve I/O or non-deterministic behavior, you can explicitly decorate them with `@task` from Prefect.
+Event stream handlers are **not automatically wrapped** by Prefect. If they involve I/O or non-deterministic behavior, you can explicitly decorate them with `@task` from Prefect. For examples, see the [streaming docs](../agents.md#streaming-all-events)
 
 The original agent, model, and MCP server can still be used as normal outside the Prefect flow.
 

pydantic_ai_slim/pydantic_ai/durable_exec/prefect/_agent.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterable, AsyncIterator, Iterator, Sequence
+from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from typing import Any, overload
@@ -50,6 +50,10 @@ def __init__(
         model_task_config: TaskConfig | None = None,
         tool_task_config: TaskConfig | None = None,
         tool_task_config_by_name: dict[str, TaskConfig | None] | None = None,
+        prefectify_toolset_func: Callable[
+            [AbstractToolset[AgentDepsT], TaskConfig, TaskConfig, dict[str, TaskConfig | None]],
+            AbstractToolset[AgentDepsT],
+        ] = prefectify_toolset,
     ):
         """Wrap an agent to enable it with Prefect durable flows, by automatically offloading model requests, tool calls, and MCP server communication to Prefect tasks.
 
@@ -63,6 +67,9 @@ def __init__(
             model_task_config: The Prefect task config to use for model request tasks. If no config is provided, use the default settings of Prefect.
             tool_task_config: The default Prefect task config to use for tool calls. If no config is provided, use the default settings of Prefect.
             tool_task_config_by_name: Per-tool task configuration. Keys are tool names, values are TaskConfig or None (None disables task wrapping for that tool).
+            prefectify_toolset_func: Optional function to use to prepare toolsets for Prefect by wrapping them in a `PrefectWrapperToolset` that moves methods that require IO to Prefect tasks.
+                If not provided, only `FunctionToolset` and `MCPServer` will be prepared for Prefect.
+                The function takes the toolset, the task config, the tool-specific task config, and the tool-specific task config by name.
         """
         super().__init__(wrapped)
 
@@ -91,23 +98,23 @@ def __init__(
         )
         self._model = prefect_model
 
-        prefect_toolsets = [toolset.visit_and_replace(self._prefectify_toolset) for toolset in wrapped.toolsets]
+        def _prefectify_toolset(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
+            """Convert a toolset to its Prefect equivalent."""
+            return prefectify_toolset_func(
+                toolset,
+                self._mcp_task_config,
+                self._tool_task_config,
+                self._tool_task_config_by_name,
+            )
+
+        prefect_toolsets = [toolset.visit_and_replace(_prefectify_toolset) for toolset in wrapped.toolsets]
         self._toolsets = prefect_toolsets
 
         # Context variable to track when we're inside this agent's Prefect flow
         self._in_prefect_agent_flow: ContextVar[bool] = ContextVar(
             f'_in_prefect_agent_flow_{self._name}', default=False
         )
 
-    def _prefectify_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
-        """Convert a toolset to its Prefect equivalent."""
-        return prefectify_toolset(
-            toolset,
-            mcp_task_config=self._mcp_task_config,
-            tool_task_config=self._tool_task_config,
-            tool_task_config_by_name=self._tool_task_config_by_name,
-        )
-
     @property
     def name(self) -> str | None:
         return self._name
@@ -812,7 +819,7 @@ def override(
             tools: The tools to use instead of the tools registered with the agent.
             instructions: The instructions to use instead of the instructions registered with the agent.
         """
-        if _utils.is_set(model) and not isinstance(model, (PrefectModel)):
+        if _utils.is_set(model) and not isinstance(model, PrefectModel):
             raise UserError(
                 'Non-Prefect model cannot be contextually overridden inside a Prefect flow, it must be set at agent creation time.'
             )

pydantic_ai_slim/pydantic_ai/durable_exec/prefect/_model.py

@@ -117,12 +117,9 @@ async def request(
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         """Make a model request, wrapped as a Prefect task when in a flow."""
-        # Get model name for task description
-        model_name = getattr(self.wrapped, 'model_name', 'unknown')
-
-        return await self._wrapped_request.with_options(name=f'Model Request: {model_name}', **self.task_config)(
-            messages, model_settings, model_request_parameters
-        )
+        return await self._wrapped_request.with_options(
+            name=f'Model Request: {self.wrapped.model_name}', **self.task_config
+        )(messages, model_settings, model_request_parameters)
 
     @asynccontextmanager
     async def request_stream(
@@ -149,10 +146,7 @@ async def request_stream(
                 return
 
         # If in a flow, consume the stream in a task and return the final response
-        # Get model name for task description
-        model_name = getattr(self.wrapped, 'model_name', 'unknown')
-
         response = await self._wrapped_request_stream.with_options(
-            name=f'Model Request (Streaming): {model_name}', **self.task_config
+            name=f'Model Request (Streaming): {self.wrapped.model_name}', **self.task_config
         )(messages, model_settings, model_request_parameters, run_context)
         yield PrefectStreamedResponse(model_request_parameters, response)