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

pydantic_ai/mcp.py

@@ -7,7 +7,7 @@ from abc import ABC, abstractmethod
 from asyncio import Lock
 from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
-from dataclasses import dataclass, field, replace
+from dataclasses import field, replace
 from datetime import timedelta
 from pathlib import Path
 from typing import Any, Callable
@@ -56,17 +56,17 @@ class MCPServer(AbstractToolset[Any], ABC):
     See <https://modelcontextprotocol.io> for more information.
     """
 
-    # these fields should be re-defined by dataclass subclasses so they appear as fields {
-    tool_prefix: str | None = None
-    log_level: mcp_types.LoggingLevel | None = None
-    log_handler: LoggingFnT | None = None
-    timeout: float = 5
-    read_timeout: float = 5 * 60
-    process_tool_call: ProcessToolCallback | None = None
-    allow_sampling: bool = True
-    max_retries: int = 1
-    sampling_model: models.Model | None = None
-    # } end of "abstract fields"
+    tool_prefix: str | None
+    log_level: mcp_types.LoggingLevel | None
+    log_handler: LoggingFnT | None
+    timeout: float
+    read_timeout: float
+    process_tool_call: ProcessToolCallback | None
+    allow_sampling: bool
+    sampling_model: models.Model | None
+    max_retries: int
+
+    _id: str | None
 
     _enter_lock: Lock = field(compare=False)
     _running_count: int
@@ -76,6 +76,34 @@ class MCPServer(AbstractToolset[Any], ABC):
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
 
+    def __init__(
+        self,
+        tool_prefix: str | None = None,
+        log_level: mcp_types.LoggingLevel | None = None,
+        log_handler: LoggingFnT | None = None,
+        timeout: float = 5,
+        read_timeout: float = 5 * 60,
+        process_tool_call: ProcessToolCallback | None = None,
+        allow_sampling: bool = True,
+        sampling_model: models.Model | None = None,
+        max_retries: int = 1,
+        *,
+        id: str | None = None,
+    ):
+        self.tool_prefix = tool_prefix
+        self.log_level = log_level
+        self.log_handler = log_handler
+        self.timeout = timeout
+        self.read_timeout = read_timeout
+        self.process_tool_call = process_tool_call
+        self.allow_sampling = allow_sampling
+        self.sampling_model = sampling_model
+        self.max_retries = max_retries
+
+        self._id = id or tool_prefix
+
+        self.__post_init__()
+
     def __post_init__(self):
         self._enter_lock = Lock()
         self._running_count = 0
@@ -96,12 +124,19 @@ class MCPServer(AbstractToolset[Any], ABC):
         yield
 
     @property
-    def name(self) -> str:
-        return repr(self)
+    def id(self) -> str | None:
+        return self._id
+
+    @property
+    def label(self) -> str:
+        if self.id:
+            return super().label  # pragma: no cover
+        else:
+            return repr(self)
 
     @property
     def tool_name_conflict_hint(self) -> str:
-        return 'Consider setting `tool_prefix` to avoid name conflicts.'
+        return 'Set the `tool_prefix` attribute to avoid name conflicts.'
 
     async def list_tools(self) -> list[mcp_types.Tool]:
         """Retrieve tools that are currently active on the server.
@@ -177,20 +212,25 @@ class MCPServer(AbstractToolset[Any], ABC):
 
     async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
         return {
-            name: ToolsetTool(
-                toolset=self,
-                tool_def=ToolDefinition(
+            name: self.tool_for_tool_def(
+                ToolDefinition(
                     name=name,
                     description=mcp_tool.description,
                     parameters_json_schema=mcp_tool.inputSchema,
                 ),
-                max_retries=self.max_retries,
-                args_validator=TOOL_SCHEMA_VALIDATOR,
             )
             for mcp_tool in await self.list_tools()
             if (name := f'{self.tool_prefix}_{mcp_tool.name}' if self.tool_prefix else mcp_tool.name)
         }
 
+    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[Any]:
+        return ToolsetTool(
+            toolset=self,
+            tool_def=tool_def,
+            max_retries=self.max_retries,
+            args_validator=TOOL_SCHEMA_VALIDATOR,
+        )
+
     async def __aenter__(self) -> Self:
         """Enter the MCP server context.
 
@@ -308,7 +348,6 @@ class MCPServer(AbstractToolset[Any], ABC):
             assert_never(resource)
 
 
-@dataclass
 class MCPServerStdio(MCPServer):
     """Runs an MCP server in a subprocess and communicates with it over stdin/stdout.
 
@@ -353,18 +392,18 @@ class MCPServerStdio(MCPServer):
     args: Sequence[str]
     """The arguments to pass to the command."""
 
-    env: dict[str, str] | None = None
+    env: dict[str, str] | None
     """The environment variables the CLI server will have access to.
 
     By default the subprocess will not inherit any environment variables from the parent process.
     If you want to inherit the environment variables from the parent process, use `env=os.environ`.
     """
 
-    cwd: str | Path | None = None
+    cwd: str | Path | None
     """The working directory to use when spawning the process."""
 
     # last fields are re-defined from the parent class so they appear as fields
-    tool_prefix: str | None = None
+    tool_prefix: str | None
     """A prefix to add to all tools that are registered with the server.
 
     If not empty, will include a trailing underscore(`_`).
@@ -372,7 +411,7 @@ class MCPServerStdio(MCPServer):
     e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
-    log_level: mcp_types.LoggingLevel | None = None
+    log_level: mcp_types.LoggingLevel | None
     """The log level to set when connecting to the server, if any.
 
     See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
@@ -380,23 +419,85 @@ class MCPServerStdio(MCPServer):
     If `None`, no log level will be set.
     """
 
-    log_handler: LoggingFnT | None = None
+    log_handler: LoggingFnT | None
     """A handler for logging messages from the server."""
 
-    timeout: float = 5
+    timeout: float
     """The timeout in seconds to wait for the client to initialize."""
 
-    process_tool_call: ProcessToolCallback | None = None
+    read_timeout: float
+    """Maximum time in seconds to wait for new messages before timing out.
+
+    This timeout applies to the long-lived connection after it's established.
+    If no new messages are received within this time, the connection will be considered stale
+    and may be closed. Defaults to 5 minutes (300 seconds).
+    """
+
+    process_tool_call: ProcessToolCallback | None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
-    allow_sampling: bool = True
+    allow_sampling: bool
     """Whether to allow MCP sampling through this client."""
 
-    max_retries: int = 1
+    sampling_model: models.Model | None
+    """The model to use for sampling."""
+
+    max_retries: int
     """The maximum number of times to retry a tool call."""
 
-    sampling_model: models.Model | None = None
-    """The model to use for sampling."""
+    def __init__(
+        self,
+        command: str,
+        args: Sequence[str],
+        env: dict[str, str] | None = None,
+        cwd: str | Path | None = None,
+        tool_prefix: str | None = None,
+        log_level: mcp_types.LoggingLevel | None = None,
+        log_handler: LoggingFnT | None = None,
+        timeout: float = 5,
+        read_timeout: float = 5 * 60,
+        process_tool_call: ProcessToolCallback | None = None,
+        allow_sampling: bool = True,
+        sampling_model: models.Model | None = None,
+        max_retries: int = 1,
+        *,
+        id: str | None = None,
+    ):
+        """Build a new MCP server.
+
+        Args:
+            command: The command to run.
+            args: The arguments to pass to the command.
+            env: The environment variables to set in the subprocess.
+            cwd: The working directory to use when spawning the process.
+            tool_prefix: A prefix to add to all tools that are registered with the server.
+            log_level: The log level to set when connecting to the server, if any.
+            log_handler: A handler for logging messages from the server.
+            timeout: The timeout in seconds to wait for the client to initialize.
+            read_timeout: Maximum time in seconds to wait for new messages before timing out.
+            process_tool_call: Hook to customize tool calling and optionally pass extra metadata.
+            allow_sampling: Whether to allow MCP sampling through this client.
+            sampling_model: The model to use for sampling.
+            max_retries: The maximum number of times to retry a tool call.
+            id: An optional unique ID for the MCP server. An MCP server 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 server's activities within the workflow.
+        """
+        self.command = command
+        self.args = args
+        self.env = env
+        self.cwd = cwd
+
+        super().__init__(
+            tool_prefix,
+            log_level,
+            log_handler,
+            timeout,
+            read_timeout,
+            process_tool_call,
+            allow_sampling,
+            sampling_model,
+            max_retries,
+            id=id,
+        )
 
     @asynccontextmanager
     async def client_streams(
@@ -412,15 +513,20 @@ class MCPServerStdio(MCPServer):
             yield read_stream, write_stream
 
     def __repr__(self) -> str:
-        return f'MCPServerStdio(command={self.command!r}, args={self.args!r}, tool_prefix={self.tool_prefix!r})'
+        repr_args = [
+            f'command={self.command!r}',
+            f'args={self.args!r}',
+        ]
+        if self.id:
+            repr_args.append(f'id={self.id!r}')  # pragma: no cover
+        return f'{self.__class__.__name__}({", ".join(repr_args)})'
 
 
-@dataclass(init=False)
 class _MCPServerHTTP(MCPServer):
     url: str
     """The URL of the endpoint on the MCP server."""
 
-    headers: dict[str, Any] | None = None
+    headers: dict[str, Any] | None
     """Optional HTTP headers to be sent with each request to the endpoint.
 
     These headers will be passed directly to the underlying `httpx.AsyncClient`.
@@ -432,7 +538,7 @@ class _MCPServerHTTP(MCPServer):
         See [`MCPServerHTTP.http_client`][pydantic_ai.mcp.MCPServerHTTP.http_client] for more information.
     """
 
-    http_client: httpx.AsyncClient | None = None
+    http_client: httpx.AsyncClient | None
     """An `httpx.AsyncClient` to use with the endpoint.
 
     This client may be configured to use customized connection parameters like self-signed certificates.
@@ -452,16 +558,8 @@ class _MCPServerHTTP(MCPServer):
         ```
     """
 
-    read_timeout: float = 5 * 60
-    """Maximum time in seconds to wait for new messages before timing out.
-
-    This timeout applies to the long-lived connection after it's established.
-    If no new messages are received within this time, the connection will be considered stale
-    and may be closed. Defaults to 5 minutes (300 seconds).
-    """
-
     # last fields are re-defined from the parent class so they appear as fields
-    tool_prefix: str | None = None
+    tool_prefix: str | None
     """A prefix to add to all tools that are registered with the server.
 
     If not empty, will include a trailing underscore (`_`).
@@ -469,7 +567,7 @@ class _MCPServerHTTP(MCPServer):
     For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
-    log_level: mcp_types.LoggingLevel | None = None
+    log_level: mcp_types.LoggingLevel | None
     """The log level to set when connecting to the server, if any.
 
     See <https://modelcontextprotocol.io/introduction#logging> for more details.
@@ -477,56 +575,81 @@ class _MCPServerHTTP(MCPServer):
     If `None`, no log level will be set.
     """
 
-    log_handler: LoggingFnT | None = None
+    log_handler: LoggingFnT | None
     """A handler for logging messages from the server."""
 
-    timeout: float = 5
+    timeout: float
     """Initial connection timeout in seconds for establishing the connection.
 
     This timeout applies to the initial connection setup and handshake.
     If the connection cannot be established within this time, the operation will fail.
     """
 
-    process_tool_call: ProcessToolCallback | None = None
+    read_timeout: float
+    """Maximum time in seconds to wait for new messages before timing out.
+
+    This timeout applies to the long-lived connection after it's established.
+    If no new messages are received within this time, the connection will be considered stale
+    and may be closed. Defaults to 5 minutes (300 seconds).
+    """
+
+    process_tool_call: ProcessToolCallback | None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
-    allow_sampling: bool = True
+    allow_sampling: bool
     """Whether to allow MCP sampling through this client."""
 
-    max_retries: int = 1
-    """The maximum number of times to retry a tool call."""
-
-    sampling_model: models.Model | None = None
+    sampling_model: models.Model | None
     """The model to use for sampling."""
 
+    max_retries: int
+    """The maximum number of times to retry a tool call."""
+
     def __init__(
         self,
         *,
         url: str,
         headers: dict[str, str] | None = None,
         http_client: httpx.AsyncClient | None = None,
-        read_timeout: float | None = None,
+        id: str | None = None,
         tool_prefix: str | None = None,
         log_level: mcp_types.LoggingLevel | None = None,
         log_handler: LoggingFnT | None = None,
         timeout: float = 5,
+        read_timeout: float | None = None,
         process_tool_call: ProcessToolCallback | None = None,
         allow_sampling: bool = True,
-        max_retries: int = 1,
         sampling_model: models.Model | None = None,
-        **kwargs: Any,
+        max_retries: int = 1,
+        **_deprecated_kwargs: Any,
     ):
-        # Handle deprecated sse_read_timeout parameter
-        if 'sse_read_timeout' in kwargs:
+        """Build a new MCP server.
+
+        Args:
+            url: The URL of the endpoint on the MCP server.
+            headers: Optional HTTP headers to be sent with each request to the endpoint.
+            http_client: An `httpx.AsyncClient` to use with the endpoint.
+            id: An optional unique ID for the MCP server. An MCP server 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 server's activities within the workflow.
+            tool_prefix: A prefix to add to all tools that are registered with the server.
+            log_level: The log level to set when connecting to the server, if any.
+            log_handler: A handler for logging messages from the server.
+            timeout: The timeout in seconds to wait for the client to initialize.
+            read_timeout: Maximum time in seconds to wait for new messages before timing out.
+            process_tool_call: Hook to customize tool calling and optionally pass extra metadata.
+            allow_sampling: Whether to allow MCP sampling through this client.
+            sampling_model: The model to use for sampling.
+            max_retries: The maximum number of times to retry a tool call.
+        """
+        if 'sse_read_timeout' in _deprecated_kwargs:
             if read_timeout is not None:
                 raise TypeError("'read_timeout' and 'sse_read_timeout' cannot be set at the same time.")
 
             warnings.warn(
                 "'sse_read_timeout' is deprecated, use 'read_timeout' instead.", DeprecationWarning, stacklevel=2
             )
-            read_timeout = kwargs.pop('sse_read_timeout')
+            read_timeout = _deprecated_kwargs.pop('sse_read_timeout')
 
-        _utils.validate_empty_kwargs(kwargs)
+        _utils.validate_empty_kwargs(_deprecated_kwargs)
 
         if read_timeout is None:
             read_timeout = 5 * 60
@@ -534,16 +657,19 @@ class _MCPServerHTTP(MCPServer):
         self.url = url
         self.headers = headers
         self.http_client = http_client
-        self.tool_prefix = tool_prefix
-        self.log_level = log_level
-        self.log_handler = log_handler
-        self.timeout = timeout
-        self.process_tool_call = process_tool_call
-        self.allow_sampling = allow_sampling
-        self.max_retries = max_retries
-        self.sampling_model = sampling_model
-        self.read_timeout = read_timeout
-        self.__post_init__()
+
+        super().__init__(
+            tool_prefix,
+            log_level,
+            log_handler,
+            timeout,
+            read_timeout,
+            process_tool_call,
+            allow_sampling,
+            sampling_model,
+            max_retries,
+            id=id,
+        )
 
     @property
     @abstractmethod
@@ -606,10 +732,14 @@ class _MCPServerHTTP(MCPServer):
                 yield read_stream, write_stream
 
     def __repr__(self) -> str:  # pragma: no cover
-        return f'{self.__class__.__name__}(url={self.url!r}, tool_prefix={self.tool_prefix!r})'
+        repr_args = [
+            f'url={self.url!r}',
+        ]
+        if self.id:
+            repr_args.append(f'id={self.id!r}')
+        return f'{self.__class__.__name__}({", ".join(repr_args)})'
 
 
-@dataclass(init=False)
 class MCPServerSSE(_MCPServerHTTP):
     """An MCP server that connects over streamable HTTP connections.
 
@@ -643,7 +773,6 @@ class MCPServerSSE(_MCPServerHTTP):
 
 
 @deprecated('The `MCPServerHTTP` class is deprecated, use `MCPServerSSE` instead.')
-@dataclass
 class MCPServerHTTP(MCPServerSSE):
     """An MCP server that connects over HTTP using the old SSE transport.
 
@@ -672,7 +801,6 @@ class MCPServerHTTP(MCPServerSSE):
     """
 
 
-@dataclass
 class MCPServerStreamableHTTP(_MCPServerHTTP):
     """An MCP server that connects over HTTP using the Streamable HTTP transport.
 

pydantic_ai/messages.py

@@ -490,8 +490,8 @@ _video_format_lookup: dict[str, VideoFormat] = {
 class UserPromptPart:
     """A user prompt, generally written by the end user.
 
-    Content comes from the `user_prompt` parameter of [`Agent.run`][pydantic_ai.Agent.run],
-    [`Agent.run_sync`][pydantic_ai.Agent.run_sync], and [`Agent.run_stream`][pydantic_ai.Agent.run_stream].
+    Content comes from the `user_prompt` parameter of [`Agent.run`][pydantic_ai.agent.AbstractAgent.run],
+    [`Agent.run_sync`][pydantic_ai.agent.AbstractAgent.run_sync], and [`Agent.run_stream`][pydantic_ai.agent.AbstractAgent.run_stream].
     """
 
     content: str | Sequence[UserContent]
@@ -532,8 +532,8 @@ tool_return_ta: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter(
 
 
 @dataclass(repr=False)
-class ToolReturnPart:
-    """A tool return message, this encodes the result of running a tool."""
+class BaseToolReturnPart:
+    """Base class for tool return parts."""
 
     tool_name: str
     """The name of the "tool" was called."""
@@ -550,9 +550,6 @@ class ToolReturnPart:
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp, when the tool returned."""
 
-    part_kind: Literal['tool-return'] = 'tool-return'
-    """Part type identifier, this is available on all parts as a discriminator."""
-
     def model_response_str(self) -> str:
         """Return a string representation of the content for the model."""
         if isinstance(self.content, str):
@@ -579,9 +576,32 @@ class ToolReturnPart:
             },
         )
 
+    def has_content(self) -> bool:
+        """Return `True` if the tool return has content."""
+        return self.content is not None  # pragma: no cover
+
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+@dataclass(repr=False)
+class ToolReturnPart(BaseToolReturnPart):
+    """A tool return message, this encodes the result of running a tool."""
+
+    part_kind: Literal['tool-return'] = 'tool-return'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
+@dataclass(repr=False)
+class BuiltinToolReturnPart(BaseToolReturnPart):
+    """A tool return message from a built-in tool."""
+
+    provider_name: str | None = None
+    """The name of the provider that generated the response."""
+
+    part_kind: Literal['builtin-tool-return'] = 'builtin-tool-return'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
 error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
 
 
@@ -723,7 +743,7 @@ class ThinkingPart:
 
 
 @dataclass(repr=False)
-class ToolCallPart:
+class BaseToolCallPart:
     """A tool call from a model."""
 
     tool_name: str
@@ -741,9 +761,6 @@ class ToolCallPart:
     In case the tool call id is not provided by the model, Pydantic AI will generate a random one.
     """
 
-    part_kind: Literal['tool-call'] = 'tool-call'
-    """Part type identifier, this is available on all parts as a discriminator."""
-
     def args_as_dict(self) -> dict[str, Any]:
         """Return the arguments as a Python dictionary.
 
@@ -780,7 +797,29 @@ class ToolCallPart:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-ModelResponsePart = Annotated[Union[TextPart, ToolCallPart, ThinkingPart], pydantic.Discriminator('part_kind')]
+@dataclass(repr=False)
+class ToolCallPart(BaseToolCallPart):
+    """A tool call from a model."""
+
+    part_kind: Literal['tool-call'] = 'tool-call'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
+@dataclass(repr=False)
+class BuiltinToolCallPart(BaseToolCallPart):
+    """A tool call to a built-in tool."""
+
+    provider_name: str | None = None
+    """The name of the provider that generated the response."""
+
+    part_kind: Literal['builtin-tool-call'] = 'builtin-tool-call'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
+ModelResponsePart = Annotated[
+    Union[TextPart, ToolCallPart, BuiltinToolCallPart, BuiltinToolReturnPart, ThinkingPart],
+    pydantic.Discriminator('part_kind'),
+]
 """A message part returned by a model."""
 
 
@@ -1172,6 +1211,29 @@ class FunctionToolResultEvent:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+@dataclass(repr=False)
+class BuiltinToolCallEvent:
+    """An event indicating the start to a call to a built-in tool."""
+
+    part: BuiltinToolCallPart
+    """The built-in tool call to make."""
+
+    event_kind: Literal['builtin_tool_call'] = 'builtin_tool_call'
+    """Event type identifier, used as a discriminator."""
+
+
+@dataclass(repr=False)
+class BuiltinToolResultEvent:
+    """An event indicating the result of a built-in tool call."""
+
+    result: BuiltinToolReturnPart
+    """The result of the call to the built-in tool."""
+
+    event_kind: Literal['builtin_tool_result'] = 'builtin_tool_result'
+    """Event type identifier, used as a discriminator."""
+
+
 HandleResponseEvent = Annotated[
-    Union[FunctionToolCallEvent, FunctionToolResultEvent], pydantic.Discriminator('event_kind')
+    Union[FunctionToolCallEvent, FunctionToolResultEvent, BuiltinToolCallEvent, BuiltinToolResultEvent],
+    pydantic.Discriminator('event_kind'),
 ]