pydantic-ai-slim 0.6.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

pydantic_ai/result.py

@@ -1,6 +1,6 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from copy import copy
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -22,7 +22,7 @@ from ._output import (
     ToolOutputSchema,
 )
 from ._run_context import AgentDepsT, RunContext
-from .messages import AgentStreamEvent, FinalResultEvent
+from .messages import AgentStreamEvent
 from .output import (
     OutputDataT,
     ToolOutput,
@@ -45,13 +45,13 @@ T = TypeVar('T')
 class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _raw_stream_response: models.StreamedResponse
     _output_schema: OutputSchema[OutputDataT]
+    _model_request_parameters: models.ModelRequestParameters
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _run_ctx: RunContext[AgentDepsT]
     _usage_limits: UsageLimits | None
     _tool_manager: ToolManager[AgentDepsT]
 
     _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
-    _final_result_event: FinalResultEvent | None = field(default=None, init=False)
     _initial_run_ctx_usage: Usage = field(init=False)
 
     def __post_init__(self):
@@ -60,12 +60,12 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
     async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[OutputDataT]:
         """Asynchronously stream the (validated) agent outputs."""
         async for response in self.stream_responses(debounce_by=debounce_by):
-            if self._final_result_event is not None:
+            if self._raw_stream_response.final_result_event is not None:
                 try:
                     yield await self._validate_response(response, allow_partial=True)
                 except ValidationError:
                     pass
-        if self._final_result_event is not None:  # pragma: no branch
+        if self._raw_stream_response.final_result_event is not None:  # pragma: no branch
             yield await self._validate_response(self._raw_stream_response.get())
 
     async def stream_responses(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
@@ -131,10 +131,11 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
 
     async def _validate_response(self, message: _messages.ModelResponse, *, allow_partial: bool = False) -> OutputDataT:
         """Validate a structured result message."""
-        if self._final_result_event is None:
+        final_result_event = self._raw_stream_response.final_result_event
+        if final_result_event is None:
             raise exceptions.UnexpectedModelBehavior('Invalid response, unable to find output')  # pragma: no cover
 
-        output_tool_name = self._final_result_event.tool_name
+        output_tool_name = final_result_event.tool_name
 
         if isinstance(self._output_schema, ToolOutputSchema) and output_tool_name is not None:
             tool_call = next(
@@ -195,7 +196,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                     and isinstance(event.part, _messages.TextPart)
                     and event.part.content
                 ):
-                    yield event.part.content, event.index  # pragma: no cover
+                    yield event.part.content, event.index
                 elif (  # pragma: no branch
                     isinstance(event, _messages.PartDeltaEvent)
                     and isinstance(event.delta, _messages.TextPartDelta)
@@ -221,52 +222,12 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 yield ''.join(deltas)
 
     def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
-        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
-
-        This proxies the _raw_stream_response and sends all events to the agent stream, while also checking for matches
-        on the result schema and emitting a [`FinalResultEvent`][pydantic_ai.messages.FinalResultEvent] if/when the
-        first match is found.
-        """
-        if self._agent_stream_iterator is not None:
-            return self._agent_stream_iterator
-
-        async def aiter():
-            output_schema = self._output_schema
-
-            def _get_final_result_event(e: _messages.ModelResponseStreamEvent) -> _messages.FinalResultEvent | None:
-                """Return an appropriate FinalResultEvent if `e` corresponds to a part that will produce a final result."""
-                if isinstance(e, _messages.PartStartEvent):
-                    new_part = e.part
-                    if isinstance(new_part, _messages.TextPart) and isinstance(
-                        output_schema, TextOutputSchema
-                    ):  # pragma: no branch
-                        return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
-                    elif isinstance(new_part, _messages.ToolCallPart) and (
-                        tool_def := self._tool_manager.get_tool_def(new_part.tool_name)
-                    ):
-                        if tool_def.kind == 'output':
-                            return _messages.FinalResultEvent(
-                                tool_name=new_part.tool_name, tool_call_id=new_part.tool_call_id
-                            )
-                        elif tool_def.kind == 'deferred':
-                            return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
-
-            usage_checking_stream = _get_usage_checking_stream_response(
+        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s."""
+        if self._agent_stream_iterator is None:
+            self._agent_stream_iterator = _get_usage_checking_stream_response(
                 self._raw_stream_response, self._usage_limits, self.usage
             )
-            async for event in usage_checking_stream:
-                yield event
-                if (final_result_event := _get_final_result_event(event)) is not None:
-                    self._final_result_event = final_result_event
-                    yield final_result_event
-                    break
-
-            # If we broke out of the above loop, we need to yield the rest of the events
-            # If we didn't, this will just be a no-op
-            async for event in usage_checking_stream:
-                yield event
-
-        self._agent_stream_iterator = aiter()
+
         return self._agent_stream_iterator
 
 
@@ -462,10 +423,10 @@ class FinalResult(Generic[OutputDataT]):
 
 
 def _get_usage_checking_stream_response(
-    stream_response: AsyncIterable[_messages.ModelResponseStreamEvent],
+    stream_response: models.StreamedResponse,
     limits: UsageLimits | None,
     get_usage: Callable[[], Usage],
-) -> AsyncIterable[_messages.ModelResponseStreamEvent]:
+) -> AsyncIterator[AgentStreamEvent]:
     if limits is not None and limits.has_token_limits():
 
         async def _usage_checking_iterator():
@@ -475,4 +436,9 @@ def _get_usage_checking_stream_response(
 
         return _usage_checking_iterator()
     else:
-        return stream_response
+        # TODO: Use `return aiter(stream_response)` once we drop support for Python 3.9
+        async def _iterator():
+            async for item in stream_response:
+                yield item
+
+        return _iterator()

pydantic_ai/run.py

@@ -0,0 +1,357 @@
+from __future__ import annotations as _annotations
+
+import dataclasses
+from collections.abc import AsyncIterator
+from copy import deepcopy
+from typing import Any, Generic, overload
+
+from typing_extensions import Literal
+
+from pydantic_graph import End, GraphRun, GraphRunContext
+
+from . import (
+    _agent_graph,
+    exceptions,
+    messages as _messages,
+    usage as _usage,
+)
+from .output import OutputDataT
+from .result import FinalResult
+from .tools import AgentDepsT
+
+
+@dataclasses.dataclass(repr=False)
+class AgentRun(Generic[AgentDepsT, OutputDataT]):
+    """A stateful, async-iterable run of an [`Agent`][pydantic_ai.agent.Agent].
+
+    You generally obtain an `AgentRun` instance by calling `async with my_agent.iter(...) as agent_run:`.
+
+    Once you have an instance, you can use it to iterate through the run's nodes as they execute. When an
+    [`End`][pydantic_graph.nodes.End] is reached, the run finishes and [`result`][pydantic_ai.agent.AgentRun.result]
+    becomes available.
+
+    Example:
+    ```python
+    from pydantic_ai import Agent
+
+    agent = Agent('openai:gpt-4o')
+
+    async def main():
+        nodes = []
+        # Iterate through the run, recording each node along the way:
+        async with agent.iter('What is the capital of France?') as agent_run:
+            async for node in agent_run:
+                nodes.append(node)
+        print(nodes)
+        '''
+        [
+            UserPromptNode(
+                user_prompt='What is the capital of France?',
+                instructions=None,
+                instructions_functions=[],
+                system_prompts=(),
+                system_prompt_functions=[],
+                system_prompt_dynamic_functions={},
+            ),
+            ModelRequestNode(
+                request=ModelRequest(
+                    parts=[
+                        UserPromptPart(
+                            content='What is the capital of France?',
+                            timestamp=datetime.datetime(...),
+                        )
+                    ]
+                )
+            ),
+            CallToolsNode(
+                model_response=ModelResponse(
+                    parts=[TextPart(content='The capital of France is Paris.')],
+                    usage=Usage(
+                        requests=1, request_tokens=56, response_tokens=7, total_tokens=63
+                    ),
+                    model_name='gpt-4o',
+                    timestamp=datetime.datetime(...),
+                )
+            ),
+            End(data=FinalResult(output='The capital of France is Paris.')),
+        ]
+        '''
+        print(agent_run.result.output)
+        #> The capital of France is Paris.
+    ```
+
+    You can also manually drive the iteration using the [`next`][pydantic_ai.agent.AgentRun.next] method for
+    more granular control.
+    """
+
+    _graph_run: GraphRun[
+        _agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any], FinalResult[OutputDataT]
+    ]
+
+    @overload
+    def _traceparent(self, *, required: Literal[False]) -> str | None: ...
+    @overload
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        traceparent = self._graph_run._traceparent(required=False)  # type: ignore[reportPrivateUsage]
+        if traceparent is None and required:  # pragma: no cover
+            raise AttributeError('No span was created for this agent run')
+        return traceparent
+
+    @property
+    def ctx(self) -> GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any]]:
+        """The current context of the agent run."""
+        return GraphRunContext[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any]](
+            self._graph_run.state, self._graph_run.deps
+        )
+
+    @property
+    def next_node(
+        self,
+    ) -> _agent_graph.AgentNode[AgentDepsT, OutputDataT] | End[FinalResult[OutputDataT]]:
+        """The next node that will be run in the agent graph.
+
+        This is the next node that will be used during async iteration, or if a node is not passed to `self.next(...)`.
+        """
+        next_node = self._graph_run.next_node
+        if isinstance(next_node, End):
+            return next_node
+        if _agent_graph.is_agent_node(next_node):
+            return next_node
+        raise exceptions.AgentRunError(f'Unexpected node type: {type(next_node)}')  # pragma: no cover
+
+    @property
+    def result(self) -> AgentRunResult[OutputDataT] | None:
+        """The final result of the run if it has ended, otherwise `None`.
+
+        Once the run returns an [`End`][pydantic_graph.nodes.End] node, `result` is populated
+        with an [`AgentRunResult`][pydantic_ai.agent.AgentRunResult].
+        """
+        graph_run_result = self._graph_run.result
+        if graph_run_result is None:
+            return None
+        return AgentRunResult(
+            graph_run_result.output.output,
+            graph_run_result.output.tool_name,
+            graph_run_result.state,
+            self._graph_run.deps.new_message_index,
+            self._traceparent(required=False),
+        )
+
+    def __aiter__(
+        self,
+    ) -> AsyncIterator[_agent_graph.AgentNode[AgentDepsT, OutputDataT] | End[FinalResult[OutputDataT]]]:
+        """Provide async-iteration over the nodes in the agent run."""
+        return self
+
+    async def __anext__(
+        self,
+    ) -> _agent_graph.AgentNode[AgentDepsT, OutputDataT] | End[FinalResult[OutputDataT]]:
+        """Advance to the next node automatically based on the last returned node."""
+        next_node = await self._graph_run.__anext__()
+        if _agent_graph.is_agent_node(node=next_node):
+            return next_node
+        assert isinstance(next_node, End), f'Unexpected node type: {type(next_node)}'
+        return next_node
+
+    async def next(
+        self,
+        node: _agent_graph.AgentNode[AgentDepsT, OutputDataT],
+    ) -> _agent_graph.AgentNode[AgentDepsT, OutputDataT] | End[FinalResult[OutputDataT]]:
+        """Manually drive the agent run by passing in the node you want to run next.
+
+        This lets you inspect or mutate the node before continuing execution, or skip certain nodes
+        under dynamic conditions. The agent run should be stopped when you return an [`End`][pydantic_graph.nodes.End]
+        node.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent
+        from pydantic_graph import End
+
+        agent = Agent('openai:gpt-4o')
+
+        async def main():
+            async with agent.iter('What is the capital of France?') as agent_run:
+                next_node = agent_run.next_node  # start with the first node
+                nodes = [next_node]
+                while not isinstance(next_node, End):
+                    next_node = await agent_run.next(next_node)
+                    nodes.append(next_node)
+                # Once `next_node` is an End, we've finished:
+                print(nodes)
+                '''
+                [
+                    UserPromptNode(
+                        user_prompt='What is the capital of France?',
+                        instructions=None,
+                        instructions_functions=[],
+                        system_prompts=(),
+                        system_prompt_functions=[],
+                        system_prompt_dynamic_functions={},
+                    ),
+                    ModelRequestNode(
+                        request=ModelRequest(
+                            parts=[
+                                UserPromptPart(
+                                    content='What is the capital of France?',
+                                    timestamp=datetime.datetime(...),
+                                )
+                            ]
+                        )
+                    ),
+                    CallToolsNode(
+                        model_response=ModelResponse(
+                            parts=[TextPart(content='The capital of France is Paris.')],
+                            usage=Usage(
+                                requests=1,
+                                request_tokens=56,
+                                response_tokens=7,
+                                total_tokens=63,
+                            ),
+                            model_name='gpt-4o',
+                            timestamp=datetime.datetime(...),
+                        )
+                    ),
+                    End(data=FinalResult(output='The capital of France is Paris.')),
+                ]
+                '''
+                print('Final result:', agent_run.result.output)
+                #> Final result: The capital of France is Paris.
+        ```
+
+        Args:
+            node: The node to run next in the graph.
+
+        Returns:
+            The next node returned by the graph logic, or an [`End`][pydantic_graph.nodes.End] node if
+            the run has completed.
+        """
+        # Note: It might be nice to expose a synchronous interface for iteration, but we shouldn't do it
+        # on this class, or else IDEs won't warn you if you accidentally use `for` instead of `async for` to iterate.
+        next_node = await self._graph_run.next(node)
+        if _agent_graph.is_agent_node(next_node):
+            return next_node
+        assert isinstance(next_node, End), f'Unexpected node type: {type(next_node)}'
+        return next_node
+
+    def usage(self) -> _usage.Usage:
+        """Get usage statistics for the run so far, including token usage, model requests, and so on."""
+        return self._graph_run.state.usage
+
+    def __repr__(self) -> str:  # pragma: no cover
+        result = self._graph_run.result
+        result_repr = '<run not finished>' if result is None else repr(result.output)
+        return f'<{type(self).__name__} result={result_repr} usage={self.usage()}>'
+
+
+@dataclasses.dataclass
+class AgentRunResult(Generic[OutputDataT]):
+    """The final result of an agent run."""
+
+    output: OutputDataT
+    """The output data from the agent run."""
+
+    _output_tool_name: str | None = dataclasses.field(repr=False)
+    _state: _agent_graph.GraphAgentState = dataclasses.field(repr=False)
+    _new_message_index: int = dataclasses.field(repr=False)
+    _traceparent_value: str | None = dataclasses.field(repr=False)
+
+    @overload
+    def _traceparent(self, *, required: Literal[False]) -> str | None: ...
+    @overload
+    def _traceparent(self) -> str: ...
+    def _traceparent(self, *, required: bool = True) -> str | None:
+        if self._traceparent_value is None and required:  # pragma: no cover
+            raise AttributeError('No span was created for this agent run')
+        return self._traceparent_value
+
+    def _set_output_tool_return(self, return_content: str) -> list[_messages.ModelMessage]:
+        """Set return content for the output tool.
+
+        Useful if you want to continue the conversation and want to set the response to the output tool call.
+        """
+        if not self._output_tool_name:
+            raise ValueError('Cannot set output tool return content when the return type is `str`.')
+
+        messages = self._state.message_history
+        last_message = messages[-1]
+        for idx, part in enumerate(last_message.parts):
+            if isinstance(part, _messages.ToolReturnPart) and part.tool_name == self._output_tool_name:
+                # Only do deepcopy when we have to modify
+                copied_messages = list(messages)
+                copied_last = deepcopy(last_message)
+                copied_last.parts[idx].content = return_content  # type: ignore[misc]
+                copied_messages[-1] = copied_last
+                return copied_messages
+
+        raise LookupError(f'No tool call found with tool name {self._output_tool_name!r}.')
+
+    def all_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
+        """Return the history of _messages.
+
+        Args:
+            output_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the output tool call if you want to continue
+                the conversation and want to set the response to the output tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            List of messages.
+        """
+        if output_tool_return_content is not None:
+            return self._set_output_tool_return(output_tool_return_content)
+        else:
+            return self._state.message_history
+
+    def all_messages_json(self, *, output_tool_return_content: str | None = None) -> bytes:
+        """Return all messages from [`all_messages`][pydantic_ai.agent.AgentRunResult.all_messages] as JSON bytes.
+
+        Args:
+            output_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the output tool call if you want to continue
+                the conversation and want to set the response to the output tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            JSON bytes representing the messages.
+        """
+        return _messages.ModelMessagesTypeAdapter.dump_json(
+            self.all_messages(output_tool_return_content=output_tool_return_content)
+        )
+
+    def new_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
+        """Return new messages associated with this run.
+
+        Messages from older runs are excluded.
+
+        Args:
+            output_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the output tool call if you want to continue
+                the conversation and want to set the response to the output tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            List of new messages.
+        """
+        return self.all_messages(output_tool_return_content=output_tool_return_content)[self._new_message_index :]
+
+    def new_messages_json(self, *, output_tool_return_content: str | None = None) -> bytes:
+        """Return new messages from [`new_messages`][pydantic_ai.agent.AgentRunResult.new_messages] as JSON bytes.
+
+        Args:
+            output_tool_return_content: The return content of the tool call to set in the last message.
+                This provides a convenient way to modify the content of the output tool call if you want to continue
+                the conversation and want to set the response to the output tool call. If `None`, the last message will
+                not be modified.
+
+        Returns:
+            JSON bytes representing the new messages.
+        """
+        return _messages.ModelMessagesTypeAdapter.dump_json(
+            self.new_messages(output_tool_return_content=output_tool_return_content)
+        )
+
+    def usage(self) -> _usage.Usage:
+        """Return the usage of the whole run."""
+        return self._state.usage

pydantic_ai/tools.py

@@ -118,7 +118,6 @@ agent = Agent('openai:gpt-4o', prepare_tools=turn_on_strict_if_openai)
 Usage `ToolsPrepareFunc[AgentDepsT]`.
 """
 
-
 DocstringFormat = Literal['google', 'numpy', 'sphinx', 'auto']
 """Supported docstring formats.
 

pydantic_ai/toolsets/__init__.py

@@ -1,3 +1,4 @@
+from ._dynamic import ToolsetFunc
 from .abstract import AbstractToolset, ToolsetTool
 from .combined import CombinedToolset
 from .deferred import DeferredToolset
@@ -10,6 +11,7 @@ from .wrapper import WrapperToolset
 
 __all__ = (
     'AbstractToolset',
+    'ToolsetFunc',
     'ToolsetTool',
     'CombinedToolset',
     'DeferredToolset',

pydantic_ai/toolsets/_dynamic.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable
+from dataclasses import dataclass, replace
+from typing import Any, Callable, Union
+
+from typing_extensions import Self, TypeAlias
+
+from .._run_context import AgentDepsT, RunContext
+from .abstract import AbstractToolset, ToolsetTool
+
+ToolsetFunc: TypeAlias = Callable[
+    [RunContext[AgentDepsT]],
+    Union[AbstractToolset[AgentDepsT], None, Awaitable[Union[AbstractToolset[AgentDepsT], None]]],
+]
+"""A sync/async function which takes a run context and returns a toolset."""
+
+
+@dataclass
+class DynamicToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that dynamically builds a toolset using a function that takes the run context.
+
+    It should only be used during a single agent run as it stores the generated toolset.
+    To use it multiple times, copy it using `dataclasses.replace`.
+    """
+
+    toolset_func: ToolsetFunc[AgentDepsT]
+    per_run_step: bool = True
+
+    _toolset: AbstractToolset[AgentDepsT] | None = None
+    _run_step: int | None = None
+
+    @property
+    def id(self) -> str | None:
+        return None  # pragma: no cover
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        try:
+            if self._toolset is not None:
+                return await self._toolset.__aexit__(*args)
+        finally:
+            self._toolset = None
+            self._run_step = None
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        if self._toolset is None or (self.per_run_step and ctx.run_step != self._run_step):
+            if self._toolset is not None:
+                await self._toolset.__aexit__()
+
+            toolset = self.toolset_func(ctx)
+            if inspect.isawaitable(toolset):
+                toolset = await toolset
+
+            if toolset is not None:
+                await toolset.__aenter__()
+
+            self._toolset = toolset
+            self._run_step = ctx.run_step
+
+        if self._toolset is None:
+            return {}
+
+        return await self._toolset.get_tools(ctx)
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        assert self._toolset is not None
+        return await self._toolset.call_tool(name, tool_args, ctx, tool)
+
+    def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
+        if self._toolset is None:
+            super().apply(visitor)
+        else:
+            self._toolset.apply(visitor)
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        if self._toolset is None:
+            return super().visit_and_replace(visitor)
+        else:
+            return replace(self, _toolset=self._toolset.visit_and_replace(visitor))

pydantic_ai/toolsets/abstract.py

@@ -70,9 +70,23 @@ class AbstractToolset(ABC, Generic[AgentDepsT]):
     """
 
     @property
-    def name(self) -> str:
+    @abstractmethod
+    def id(self) -> str | None:
+        """An ID for the toolset that is unique among all toolsets registered with the same agent.
+
+        If you're implementing a concrete implementation that users can instantiate more than once, you should let them optionally pass a custom ID to the constructor and return that here.
+
+        A toolset needs to have an ID in order to be used in a durable execution environment like Temporal, in which case the ID will be used to identify the toolset's activities within the workflow.
+        """
+        raise NotImplementedError()
+
+    @property
+    def label(self) -> str:
         """The name of the toolset for use in error messages."""
-        return self.__class__.__name__.replace('Toolset', ' toolset')
+        label = self.__class__.__name__
+        if self.id:  # pragma: no branch
+            label += f' {self.id!r}'
+        return label
 
     @property
     def tool_name_conflict_hint(self) -> str:
@@ -113,9 +127,15 @@ class AbstractToolset(ABC, Generic[AgentDepsT]):
         raise NotImplementedError()
 
     def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
-        """Run a visitor function on all concrete toolsets that are not wrappers (i.e. they implement their own tool listing and calling)."""
+        """Run a visitor function on all "leaf" toolsets (i.e. those that implement their own tool listing and calling)."""
         visitor(self)
 
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        """Run a visitor function on all "leaf" toolsets (i.e. those that implement their own tool listing and calling) and replace them in the hierarchy with the result of the function."""
+        return visitor(self)
+
     def filtered(
         self, filter_func: Callable[[RunContext[AgentDepsT], ToolDefinition], bool]
     ) -> FilteredToolset[AgentDepsT]: