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

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]:

pydantic_ai/toolsets/combined.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Sequence
 from contextlib import AsyncExitStack
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import Any, Callable
 
 from typing_extensions import Self
@@ -40,6 +40,14 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
         self._entered_count = 0
         self._exit_stack = None
 
+    @property
+    def id(self) -> str | None:
+        return None  # pragma: no cover
+
+    @property
+    def label(self) -> str:
+        return f'{self.__class__.__name__}({", ".join(toolset.label for toolset in self.toolsets)})'  # pragma: no cover
+
     async def __aenter__(self) -> Self:
         async with self._enter_lock:
             if self._entered_count == 0:
@@ -63,13 +71,15 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
 
         for toolset, tools in zip(self.toolsets, toolsets_tools):
             for name, tool in tools.items():
-                if existing_tools := all_tools.get(name):
+                tool_toolset = tool.toolset
+                if existing_tool := all_tools.get(name):
+                    capitalized_toolset_label = tool_toolset.label[0].upper() + tool_toolset.label[1:]
                     raise UserError(
-                        f'{toolset.name} defines a tool whose name conflicts with existing tool from {existing_tools.toolset.name}: {name!r}. {toolset.tool_name_conflict_hint}'
+                        f'{capitalized_toolset_label} defines a tool whose name conflicts with existing tool from {existing_tool.toolset.label}: {name!r}. {toolset.tool_name_conflict_hint}'
                     )
 
                 all_tools[name] = _CombinedToolsetTool(
-                    toolset=tool.toolset,
+                    toolset=tool_toolset,
                     tool_def=tool.tool_def,
                     max_retries=tool.max_retries,
                     args_validator=tool.args_validator,
@@ -87,3 +97,8 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
     def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
         for toolset in self.toolsets:
             toolset.apply(visitor)
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        return replace(self, toolsets=[toolset.visit_and_replace(visitor) for toolset in self.toolsets])

pydantic_ai/toolsets/deferred.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, replace
+from dataclasses import replace
 from typing import Any
 
 from pydantic_core import SchemaValidator, core_schema
@@ -12,7 +12,6 @@ from .abstract import AbstractToolset, ToolsetTool
 TOOL_SCHEMA_VALIDATOR = SchemaValidator(schema=core_schema.any_schema())
 
 
-@dataclass
 class DeferredToolset(AbstractToolset[AgentDepsT]):
     """A toolset that holds deferred tools whose results will be produced outside of the Pydantic AI agent run in which they were called.
 
@@ -20,6 +19,15 @@ class DeferredToolset(AbstractToolset[AgentDepsT]):
     """
 
     tool_defs: list[ToolDefinition]
+    _id: str | None
+
+    def __init__(self, tool_defs: list[ToolDefinition], *, id: str | None = None):
+        self.tool_defs = tool_defs
+        self._id = id
+
+    @property
+    def id(self) -> str | None:
+        return self._id
 
     async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
         return {

pydantic_ai/toolsets/function.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Sequence
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, replace
 from typing import Any, Callable, overload
 
 from pydantic.json_schema import GenerateJsonSchema
@@ -20,30 +20,40 @@ from .abstract import AbstractToolset, ToolsetTool
 
 
 @dataclass
-class _FunctionToolsetTool(ToolsetTool[AgentDepsT]):
+class FunctionToolsetTool(ToolsetTool[AgentDepsT]):
     """A tool definition for a function toolset tool that keeps track of the function to call."""
 
     call_func: Callable[[dict[str, Any], RunContext[AgentDepsT]], Awaitable[Any]]
+    is_async: bool
 
 
-@dataclass(init=False)
 class FunctionToolset(AbstractToolset[AgentDepsT]):
     """A toolset that lets Python functions be used as tools.
 
     See [toolset docs](../toolsets.md#function-toolset) for more information.
     """
 
-    max_retries: int = field(default=1)
-    tools: dict[str, Tool[Any]] = field(default_factory=dict)
+    max_retries: int
+    tools: dict[str, Tool[Any]]
+    _id: str | None
 
-    def __init__(self, tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [], max_retries: int = 1):
+    def __init__(
+        self,
+        tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
+        max_retries: int = 1,
+        *,
+        id: str | None = None,
+    ):
         """Build a new function toolset.
 
         Args:
             tools: The tools to add to the toolset.
             max_retries: The maximum number of retries for each tool during a run.
+            id: An optional unique ID for the toolset. 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.
         """
         self.max_retries = max_retries
+        self._id = id
+
         self.tools = {}
         for tool in tools:
             if isinstance(tool, Tool):
@@ -51,6 +61,10 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
             else:
                 self.add_function(tool)
 
+    @property
+    def id(self) -> str | None:
+        return self._id
+
     @overload
     def tool(self, func: ToolFuncEither[AgentDepsT, ToolParams], /) -> ToolFuncEither[AgentDepsT, ToolParams]: ...
 
@@ -222,17 +236,18 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
                 else:
                     raise UserError(f'Tool name conflicts with previously renamed tool: {new_name!r}.')
 
-            tools[new_name] = _FunctionToolsetTool(
+            tools[new_name] = FunctionToolsetTool(
                 toolset=self,
                 tool_def=tool_def,
                 max_retries=tool.max_retries if tool.max_retries is not None else self.max_retries,
                 args_validator=tool.function_schema.validator,
                 call_func=tool.function_schema.call,
+                is_async=tool.function_schema.is_async,
             )
         return tools
 
     async def call_tool(
         self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
     ) -> Any:
-        assert isinstance(tool, _FunctionToolsetTool)
+        assert isinstance(tool, FunctionToolsetTool)
         return await tool.call_func(tool_args, ctx)