pydantic-ai-slim 0.3.0__tar.gz → 0.3.2__tar.gz

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/.gitignore

@@ -19,3 +19,4 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 node_modules/
 **.idea/
 .coverage*
+/test_tmp/

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.3.0
+Version: 0.3.2
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
 License-Expression: MIT
@@ -30,11 +30,11 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.3.0
+Requires-Dist: pydantic-graph==0.3.2
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
-Requires-Dist: fasta2a==0.3.0; extra == 'a2a'
+Requires-Dist: fasta2a==0.3.2; extra == 'a2a'
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.52.0; extra == 'anthropic'
 Provides-Extra: bedrock
@@ -48,7 +48,7 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.3.0; extra == 'evals'
+Requires-Dist: pydantic-evals==0.3.2; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.15.0; extra == 'google'
 Provides-Extra: groq

pydantic_ai_slim-0.3.2/pydantic_ai/_mcp.py

@@ -0,0 +1,123 @@
+import base64
+from collections.abc import Sequence
+from typing import Literal
+
+from . import exceptions, messages
+
+try:
+    from mcp import types as mcp_types
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install the `mcp` package to use the MCP server, '
+        'you can use the `mcp` optional group — `pip install "pydantic-ai-slim[mcp]"`'
+    ) from _import_error
+
+
+def map_from_mcp_params(params: mcp_types.CreateMessageRequestParams) -> list[messages.ModelMessage]:
+    """Convert from MCP create message request parameters to pydantic-ai messages."""
+    pai_messages: list[messages.ModelMessage] = []
+    request_parts: list[messages.ModelRequestPart] = []
+    if params.systemPrompt:
+        request_parts.append(messages.SystemPromptPart(content=params.systemPrompt))
+    response_parts: list[messages.ModelResponsePart] = []
+    for msg in params.messages:
+        content = msg.content
+        if msg.role == 'user':
+            # if there are any response parts, add a response message wrapping them
+            if response_parts:
+                pai_messages.append(messages.ModelResponse(parts=response_parts))
+                response_parts = []
+
+            # TODO(Marcelo): We can reuse the `_map_tool_result_part` from the mcp module here.
+            if isinstance(content, mcp_types.TextContent):
+                user_part_content: str | Sequence[messages.UserContent] = content.text
+            else:
+                # image content
+                user_part_content = [
+                    messages.BinaryContent(data=base64.b64decode(content.data), media_type=content.mimeType)
+                ]
+
+            request_parts.append(messages.UserPromptPart(content=user_part_content))
+        else:
+            # role is assistant
+            # if there are any request parts, add a request message wrapping them
+            if request_parts:
+                pai_messages.append(messages.ModelRequest(parts=request_parts))
+                request_parts = []
+
+            response_parts.append(map_from_sampling_content(content))
+
+    if response_parts:
+        pai_messages.append(messages.ModelResponse(parts=response_parts))
+    if request_parts:
+        pai_messages.append(messages.ModelRequest(parts=request_parts))
+    return pai_messages
+
+
+def map_from_pai_messages(pai_messages: list[messages.ModelMessage]) -> tuple[str, list[mcp_types.SamplingMessage]]:
+    """Convert from pydantic-ai messages to MCP sampling messages.
+
+    Returns:
+        A tuple containing the system prompt and a list of sampling messages.
+    """
+    sampling_msgs: list[mcp_types.SamplingMessage] = []
+
+    def add_msg(
+        role: Literal['user', 'assistant'],
+        content: mcp_types.TextContent | mcp_types.ImageContent | mcp_types.AudioContent,
+    ):
+        sampling_msgs.append(mcp_types.SamplingMessage(role=role, content=content))
+
+    system_prompt: list[str] = []
+    for pai_message in pai_messages:
+        if isinstance(pai_message, messages.ModelRequest):
+            if pai_message.instructions is not None:
+                system_prompt.append(pai_message.instructions)
+
+            for part in pai_message.parts:
+                if isinstance(part, messages.SystemPromptPart):
+                    system_prompt.append(part.content)
+                if isinstance(part, messages.UserPromptPart):
+                    if isinstance(part.content, str):
+                        add_msg('user', mcp_types.TextContent(type='text', text=part.content))
+                    else:
+                        for chunk in part.content:
+                            if isinstance(chunk, str):
+                                add_msg('user', mcp_types.TextContent(type='text', text=chunk))
+                            elif isinstance(chunk, messages.BinaryContent) and chunk.is_image:
+                                add_msg(
+                                    'user',
+                                    mcp_types.ImageContent(
+                                        type='image',
+                                        data=base64.b64decode(chunk.data).decode(),
+                                        mimeType=chunk.media_type,
+                                    ),
+                                )
+                            # TODO(Marcelo): Add support for audio content.
+                            else:
+                                raise NotImplementedError(f'Unsupported content type: {type(chunk)}')
+        else:
+            add_msg('assistant', map_from_model_response(pai_message))
+    return ''.join(system_prompt), sampling_msgs
+
+
+def map_from_model_response(model_response: messages.ModelResponse) -> mcp_types.TextContent:
+    """Convert from a model response to MCP text content."""
+    text_parts: list[str] = []
+    for part in model_response.parts:
+        if isinstance(part, messages.TextPart):
+            text_parts.append(part.content)
+        # TODO(Marcelo): We should ignore ThinkingPart here.
+        else:
+            raise exceptions.UnexpectedModelBehavior(f'Unexpected part type: {type(part).__name__}, expected TextPart')
+    return mcp_types.TextContent(type='text', text=''.join(text_parts))
+
+
+def map_from_sampling_content(
+    content: mcp_types.TextContent | mcp_types.ImageContent | mcp_types.AudioContent,
+) -> messages.TextPart:
+    """Convert from sampling content to a pydantic-ai text part."""
+    if isinstance(content, mcp_types.TextContent):  # pragma: no branch
+        return messages.TextPart(content=content.text)
+    else:
+        raise NotImplementedError('Image and Audio responses in sampling are not yet supported')

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/agent.py

@@ -1691,14 +1691,23 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         return isinstance(node, End)
 
     @asynccontextmanager
-    async def run_mcp_servers(self) -> AsyncIterator[None]:
+    async def run_mcp_servers(
+        self, model: models.Model | models.KnownModelName | str | None = None
+    ) -> AsyncIterator[None]:
         """Run [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] so they can be used by the agent.
 
         Returns: a context manager to start and shutdown the servers.
         """
+        try:
+            sampling_model: models.Model | None = self._get_model(model)
+        except exceptions.UserError:  # pragma: no cover
+            sampling_model = None
+
         exit_stack = AsyncExitStack()
         try:
             for mcp_server in self._mcp_servers:
+                if sampling_model is not None:  # pragma: no branch
+                    mcp_server.sampling_model = sampling_model
                 await exit_stack.enter_async_context(mcp_server)
             yield
         finally:

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/mcp.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import base64
 import functools
-import json
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
@@ -13,41 +12,28 @@ from typing import Any, Callable
 
 import anyio
 import httpx
+import pydantic_core
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
-from mcp.shared.exceptions import McpError
-from mcp.shared.message import SessionMessage
-from mcp.types import (
-    AudioContent,
-    BlobResourceContents,
-    CallToolRequest,
-    CallToolRequestParams,
-    CallToolResult,
-    ClientRequest,
-    Content,
-    EmbeddedResource,
-    ImageContent,
-    LoggingLevel,
-    RequestParams,
-    TextContent,
-    TextResourceContents,
-)
 from typing_extensions import Self, assert_never, deprecated
 
-from pydantic_ai.exceptions import ModelRetry
-from pydantic_ai.messages import BinaryContent
-from pydantic_ai.tools import RunContext, ToolDefinition
-
 try:
-    from mcp.client.session import ClientSession
+    from mcp import types as mcp_types
+    from mcp.client.session import ClientSession, LoggingFnT
     from mcp.client.sse import sse_client
     from mcp.client.stdio import StdioServerParameters, stdio_client
+    from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
+    from mcp.shared.context import RequestContext
+    from mcp.shared.exceptions import McpError
+    from mcp.shared.message import SessionMessage
 except ImportError as _import_error:
     raise ImportError(
         'Please install the `mcp` package to use the MCP server, '
         'you can use the `mcp` optional group — `pip install "pydantic-ai-slim[mcp]"`'
     ) from _import_error
 
+# after mcp imports so any import error maps to this file, not _mcp.py
+from . import _mcp, exceptions, messages, models, tools
+
 __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP', 'MCPServerSSE', 'MCPServerStreamableHTTP'
 
 
@@ -57,22 +43,22 @@ class MCPServer(ABC):
     See <https://modelcontextprotocol.io> for more information.
     """
 
-    is_running: bool = False
+    # these fields should be re-defined by dataclass subclasses so they appear as fields {
     tool_prefix: str | None = None
-    """A prefix to add to all tools that are registered with the server.
-
-    If not empty, will include a trailing underscore(`_`).
-
-    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_handler: LoggingFnT | None = None
+    timeout: float = 5
     process_tool_call: ProcessToolCallback | None = None
-    """Hook to customize tool calling and optionally pass extra metadata."""
+    allow_sampling: bool = True
+    # } end of "abstract fields"
+
+    _running_count: int = 0
 
     _client: ClientSession
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
     _exit_stack: AsyncExitStack
+    sampling_model: models.Model | None = None
 
     @abstractmethod
     @asynccontextmanager
@@ -88,14 +74,6 @@ class MCPServer(ABC):
         raise NotImplementedError('MCP Server subclasses must implement this method.')
         yield
 
-    @abstractmethod
-    def _get_log_level(self) -> LoggingLevel | None:
-        """Get the log level for the MCP server."""
-        raise NotImplementedError('MCP Server subclasses must implement this method.')
-
-    def _get_client_initialize_timeout(self) -> float:
-        return 5  # pragma: no cover
-
     def get_prefixed_tool_name(self, tool_name: str) -> str:
         """Get the tool name with prefix if `tool_prefix` is set."""
         return f'{self.tool_prefix}_{tool_name}' if self.tool_prefix else tool_name
@@ -104,21 +82,26 @@ class MCPServer(ABC):
         """Get original tool name without prefix for calling tools."""
         return tool_name.removeprefix(f'{self.tool_prefix}_') if self.tool_prefix else tool_name
 
-    async def list_tools(self) -> list[ToolDefinition]:
+    @property
+    def is_running(self) -> bool:
+        """Check if the MCP server is running."""
+        return bool(self._running_count)
+
+    async def list_tools(self) -> list[tools.ToolDefinition]:
         """Retrieve tools that are currently active on the server.
 
         Note:
         - We don't cache tools as they might change.
         - We also don't subscribe to the server to avoid complexity.
         """
-        tools = await self._client.list_tools()
+        mcp_tools = await self._client.list_tools()
         return [
-            ToolDefinition(
+            tools.ToolDefinition(
                 name=self.get_prefixed_tool_name(tool.name),
                 description=tool.description or '',
                 parameters_json_schema=tool.inputSchema,
             )
-            for tool in tools.tools
+            for tool in mcp_tools.tools
         ]
 
     async def call_tool(
@@ -143,44 +126,48 @@ class MCPServer(ABC):
         try:
             # meta param is not provided by session yet, so build and can send_request directly.
             result = await self._client.send_request(
-                ClientRequest(
-                    CallToolRequest(
+                mcp_types.ClientRequest(
+                    mcp_types.CallToolRequest(
                         method='tools/call',
-                        params=CallToolRequestParams(
+                        params=mcp_types.CallToolRequestParams(
                             name=self.get_unprefixed_tool_name(tool_name),
                             arguments=arguments,
-                            _meta=RequestParams.Meta(**metadata) if metadata else None,
+                            _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
                         ),
                     )
                 ),
-                CallToolResult,
+                mcp_types.CallToolResult,
             )
         except McpError as e:
-            raise ModelRetry(e.error.message)
+            raise exceptions.ModelRetry(e.error.message)
 
         content = [self._map_tool_result_part(part) for part in result.content]
 
         if result.isError:
             text = '\n'.join(str(part) for part in content)
-            raise ModelRetry(text)
-
-        if len(content) == 1:
-            return content[0]
-        return content
+            raise exceptions.ModelRetry(text)
+        else:
+            return content[0] if len(content) == 1 else content
 
     async def __aenter__(self) -> Self:
-        self._exit_stack = AsyncExitStack()
-
-        self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(self.client_streams())
-        client = ClientSession(read_stream=self._read_stream, write_stream=self._write_stream)
-        self._client = await self._exit_stack.enter_async_context(client)
+        if self._running_count == 0:
+            self._exit_stack = AsyncExitStack()
+
+            self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(self.client_streams())
+            client = ClientSession(
+                read_stream=self._read_stream,
+                write_stream=self._write_stream,
+                sampling_callback=self._sampling_callback if self.allow_sampling else None,
+                logging_callback=self.log_handler,
+            )
+            self._client = await self._exit_stack.enter_async_context(client)
 
-        with anyio.fail_after(self._get_client_initialize_timeout()):
-            await self._client.initialize()
+            with anyio.fail_after(self.timeout):
+                await self._client.initialize()
 
-        if log_level := self._get_log_level():
-            await self._client.set_logging_level(log_level)
-        self.is_running = True
+                if log_level := self.log_level:
+                    await self._client.set_logging_level(log_level)
+        self._running_count += 1
         return self
 
     async def __aexit__(
@@ -189,32 +176,64 @@ class MCPServer(ABC):
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> bool | None:
-        await self._exit_stack.aclose()
-        self.is_running = False
+        self._running_count -= 1
+        if self._running_count <= 0:
+            await self._exit_stack.aclose()
+
+    async def _sampling_callback(
+        self, context: RequestContext[ClientSession, Any], params: mcp_types.CreateMessageRequestParams
+    ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData:
+        """MCP sampling callback."""
+        if self.sampling_model is None:
+            raise ValueError('Sampling model is not set')  # pragma: no cover
+
+        pai_messages = _mcp.map_from_mcp_params(params)
+        model_settings = models.ModelSettings()
+        if max_tokens := params.maxTokens:  # pragma: no branch
+            model_settings['max_tokens'] = max_tokens
+        if temperature := params.temperature:  # pragma: no branch
+            model_settings['temperature'] = temperature
+        if stop_sequences := params.stopSequences:  # pragma: no branch
+            model_settings['stop_sequences'] = stop_sequences
+
+        model_response = await self.sampling_model.request(
+            pai_messages,
+            model_settings,
+            models.ModelRequestParameters(),
+        )
+        return mcp_types.CreateMessageResult(
+            role='assistant',
+            content=_mcp.map_from_model_response(model_response),
+            model=self.sampling_model.model_name,
+        )
 
-    def _map_tool_result_part(self, part: Content) -> str | BinaryContent | dict[str, Any] | list[Any]:
+    def _map_tool_result_part(
+        self, part: mcp_types.Content
+    ) -> str | messages.BinaryContent | dict[str, Any] | list[Any]:
         # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values
 
-        if isinstance(part, TextContent):
+        if isinstance(part, mcp_types.TextContent):
             text = part.text
             if text.startswith(('[', '{')):
                 try:
-                    return json.loads(text)
+                    return pydantic_core.from_json(text)
                 except ValueError:
                     pass
             return text
-        elif isinstance(part, ImageContent):
-            return BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
-        elif isinstance(part, AudioContent):
+        elif isinstance(part, mcp_types.ImageContent):
+            return messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
+        elif isinstance(part, mcp_types.AudioContent):
             # NOTE: The FastMCP server doesn't support audio content.
             # See <https://github.com/modelcontextprotocol/python-sdk/issues/952> for more details.
-            return BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)  # pragma: no cover
-        elif isinstance(part, EmbeddedResource):
+            return messages.BinaryContent(
+                data=base64.b64decode(part.data), media_type=part.mimeType
+            )  # pragma: no cover
+        elif isinstance(part, mcp_types.EmbeddedResource):
             resource = part.resource
-            if isinstance(resource, TextResourceContents):
+            if isinstance(resource, mcp_types.TextResourceContents):
                 return resource.text
-            elif isinstance(resource, BlobResourceContents):
-                return BinaryContent(
+            elif isinstance(resource, mcp_types.BlobResourceContents):
+                return messages.BinaryContent(
                     data=base64.b64decode(resource.blob),
                     media_type=resource.mimeType or 'application/octet-stream',
                 )
@@ -275,17 +294,11 @@ class MCPServerStdio(MCPServer):
     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`.
     """
-    log_level: LoggingLevel | None = 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.
-
-    If `None`, no log level will be set.
-    """
 
     cwd: str | Path | None = 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
     """A prefix to add to all tools that are registered with the server.
 
@@ -294,11 +307,25 @@ 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
+    """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.
+
+    If `None`, no log level will be set.
+    """
+
+    log_handler: LoggingFnT | None = None
+    """A handler for logging messages from the server."""
+
+    timeout: float = 5
+    """The timeout in seconds to wait for the client to initialize."""
+
     process_tool_call: ProcessToolCallback | None = None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
-    timeout: float = 5
-    """ The timeout in seconds to wait for the client to initialize."""
+    allow_sampling: bool = True
+    """Whether to allow MCP sampling through this client."""
 
     @asynccontextmanager
     async def client_streams(
@@ -313,15 +340,9 @@ class MCPServerStdio(MCPServer):
         async with stdio_client(server=server) as (read_stream, write_stream):
             yield read_stream, write_stream
 
-    def _get_log_level(self) -> LoggingLevel | None:
-        return self.log_level
-
     def __repr__(self) -> str:
         return f'MCPServerStdio(command={self.command!r}, args={self.args!r}, tool_prefix={self.tool_prefix!r})'
 
-    def _get_client_initialize_timeout(self) -> float:
-        return self.timeout
-
 
 @dataclass
 class _MCPServerHTTP(MCPServer):
@@ -360,13 +381,6 @@ class _MCPServerHTTP(MCPServer):
         ```
     """
 
-    timeout: float = 5
-    """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.
-    """
-
     sse_read_timeout: float = 5 * 60
     """Maximum time in seconds to wait for new SSE messages before timing out.
 
@@ -375,7 +389,16 @@ class _MCPServerHTTP(MCPServer):
     and may be closed. Defaults to 5 minutes (300 seconds).
     """
 
-    log_level: LoggingLevel | None = None
+    # last fields are re-defined from the parent class so they appear as fields
+    tool_prefix: str | None = None
+    """A prefix to add to all tools that are registered with the server.
+
+    If not empty, will include a trailing underscore (`_`).
+
+    For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
+    """
+
+    log_level: mcp_types.LoggingLevel | None = None
     """The log level to set when connecting to the server, if any.
 
     See <https://modelcontextprotocol.io/introduction#logging> for more details.
@@ -383,17 +406,22 @@ class _MCPServerHTTP(MCPServer):
     If `None`, no log level will be set.
     """
 
-    tool_prefix: str | None = None
-    """A prefix to add to all tools that are registered with the server.
+    log_handler: LoggingFnT | None = None
+    """A handler for logging messages from the server."""
 
-    If not empty, will include a trailing underscore (`_`).
+    timeout: float = 5
+    """Initial connection timeout in seconds for establishing the connection.
 
-    For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
+    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
     """Hook to customize tool calling and optionally pass extra metadata."""
 
+    allow_sampling: bool = True
+    """Whether to allow MCP sampling through this client."""
+
     @property
     @abstractmethod
     def _transport_client(
@@ -419,7 +447,10 @@ class _MCPServerHTTP(MCPServer):
     async def client_streams(
         self,
     ) -> AsyncIterator[
-        tuple[MemoryObjectReceiveStream[SessionMessage | Exception], MemoryObjectSendStream[SessionMessage]]
+        tuple[
+            MemoryObjectReceiveStream[SessionMessage | Exception],
+            MemoryObjectSendStream[SessionMessage],
+        ]
     ]:  # pragma: no cover
         if self.http_client and self.headers:
             raise ValueError('`http_client` is mutually exclusive with `headers`.')
@@ -451,15 +482,9 @@ class _MCPServerHTTP(MCPServer):
             async with transport_client_partial(headers=self.headers) as (read_stream, write_stream, *_):
                 yield read_stream, write_stream
 
-    def _get_log_level(self) -> LoggingLevel | None:
-        return self.log_level
-
     def __repr__(self) -> str:  # pragma: no cover
         return f'{self.__class__.__name__}(url={self.url!r}, tool_prefix={self.tool_prefix!r})'
 
-    def _get_client_initialize_timeout(self) -> float:  # pragma: no cover
-        return self.timeout
-
 
 @dataclass
 class MCPServerSSE(_MCPServerHTTP):
@@ -555,7 +580,11 @@ class MCPServerStreamableHTTP(_MCPServerHTTP):
 
 
 ToolResult = (
-    str | BinaryContent | dict[str, Any] | list[Any] | Sequence[str | BinaryContent | dict[str, Any] | list[Any]]
+    str
+    | messages.BinaryContent
+    | dict[str, Any]
+    | list[Any]
+    | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]]
 )
 """The result type of a tool call."""
 
@@ -564,7 +593,7 @@ CallToolFunc = Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[
 
 ProcessToolCallback = Callable[
     [
-        RunContext[Any],
+        tools.RunContext[Any],
         CallToolFunc,
         str,
         dict[str, Any],

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/messages.py

@@ -763,7 +763,9 @@ class ThinkingPartDelta:
             ValueError: If `part` is not a `ThinkingPart`.
         """
         if isinstance(part, ThinkingPart):
-            return replace(part, content=part.content + self.content_delta if self.content_delta else None)
+            new_content = part.content + self.content_delta if self.content_delta else part.content
+            new_signature = self.signature_delta if self.signature_delta is not None else part.signature
+            return replace(part, content=new_content, signature=new_signature)
         elif isinstance(part, ThinkingPartDelta):
             if self.content_delta is None and self.signature_delta is None:
                 raise ValueError('Cannot apply ThinkingPartDelta with no content or signature')

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/models/__init__.py

@@ -139,8 +139,11 @@ KnownModelName = TypeAliasType(
         'google-gla:gemini-2.0-flash-lite-preview-02-05',
         'google-gla:gemini-2.0-pro-exp-02-05',
         'google-gla:gemini-2.5-flash-preview-05-20',
+        'google-gla:gemini-2.5-flash',
+        'google-gla:gemini-2.5-flash-lite-preview-06-17',
         'google-gla:gemini-2.5-pro-exp-03-25',
         'google-gla:gemini-2.5-pro-preview-05-06',
+        'google-gla:gemini-2.5-pro',
         'google-vertex:gemini-1.5-flash',
         'google-vertex:gemini-1.5-flash-8b',
         'google-vertex:gemini-1.5-pro',
@@ -149,8 +152,11 @@ KnownModelName = TypeAliasType(
         'google-vertex:gemini-2.0-flash-lite-preview-02-05',
         'google-vertex:gemini-2.0-pro-exp-02-05',
         'google-vertex:gemini-2.5-flash-preview-05-20',
+        'google-vertex:gemini-2.5-flash',
+        'google-vertex:gemini-2.5-flash-lite-preview-06-17',
         'google-vertex:gemini-2.5-pro-exp-03-25',
         'google-vertex:gemini-2.5-pro-preview-05-06',
+        'google-vertex:gemini-2.5-pro',
         'gpt-3.5-turbo',
         'gpt-3.5-turbo-0125',
         'gpt-3.5-turbo-0301',

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/models/gemini.py

@@ -54,8 +54,11 @@ LatestGeminiModelNames = Literal[
     'gemini-2.0-flash-lite-preview-02-05',
     'gemini-2.0-pro-exp-02-05',
     'gemini-2.5-flash-preview-05-20',
+    'gemini-2.5-flash',
+    'gemini-2.5-flash-lite-preview-06-17',
     'gemini-2.5-pro-exp-03-25',
     'gemini-2.5-pro-preview-05-06',
+    'gemini-2.5-pro',
 ]
 """Latest Gemini models."""
 

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/models/google.py

@@ -79,8 +79,11 @@ LatestGoogleModelNames = Literal[
     'gemini-2.0-flash-lite-preview-02-05',
     'gemini-2.0-pro-exp-02-05',
     'gemini-2.5-flash-preview-05-20',
+    'gemini-2.5-flash',
+    'gemini-2.5-flash-lite-preview-06-17',
     'gemini-2.5-pro-exp-03-25',
     'gemini-2.5-pro-preview-05-06',
+    'gemini-2.5-pro',
 ]
 """Latest Gemini models."""
 

pydantic_ai_slim-0.3.2/pydantic_ai/models/mcp_sampling.py

@@ -0,0 +1,95 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, cast
+
+from .. import _mcp, exceptions, usage
+from ..messages import ModelMessage, ModelResponse
+from ..settings import ModelSettings
+from . import Model, ModelRequestParameters, StreamedResponse
+
+if TYPE_CHECKING:
+    from mcp import ServerSession
+    from mcp.types import ModelPreferences
+
+
+class MCPSamplingModelSettings(ModelSettings, total=False):
+    """Settings used for an MCP Sampling model request.
+
+    ALL FIELDS MUST BE `mcp_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+    """
+
+    mcp_model_preferences: ModelPreferences
+    """Model preferences to use for MCP Sampling."""
+
+
+@dataclass
+class MCPSamplingModel(Model):
+    """A model that uses MCP Sampling.
+
+    [MCP Sampling](https://modelcontextprotocol.io/docs/concepts/sampling)
+    allows an MCP server to make requests to a model by calling back to the MCP client that connected to it.
+    """
+
+    session: ServerSession
+    """The MCP server session to use for sampling."""
+
+    default_max_tokens: int = 16_384
+    """Default max tokens to use if not set in [`ModelSettings`][pydantic_ai.settings.ModelSettings.max_tokens].
+
+    Max tokens is a required parameter for MCP Sampling, but optional on
+    [`ModelSettings`][pydantic_ai.settings.ModelSettings], so this value is used as fallback.
+    """
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        system_prompt, sampling_messages = _mcp.map_from_pai_messages(messages)
+        model_settings = cast(MCPSamplingModelSettings, model_settings or {})
+
+        result = await self.session.create_message(
+            sampling_messages,
+            max_tokens=model_settings.get('max_tokens', self.default_max_tokens),
+            system_prompt=system_prompt,
+            temperature=model_settings.get('temperature'),
+            model_preferences=model_settings.get('mcp_model_preferences'),
+            stop_sequences=model_settings.get('stop_sequences'),
+        )
+        if result.role == 'assistant':
+            return ModelResponse(
+                parts=[_mcp.map_from_sampling_content(result.content)],
+                usage=usage.Usage(requests=1),
+                model_name=result.model,
+            )
+        else:
+            raise exceptions.UnexpectedModelBehavior(
+                f'Unexpected result from MCP sampling, expected "assistant" role, got {result.role}.'
+            )
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        raise NotImplementedError('MCP Sampling does not support streaming')
+        yield
+
+    @property
+    def model_name(self) -> str:
+        """The model name.
+
+        Since the model name isn't known until the request is made, this property always returns `'mcp-sampling'`.
+        """
+        return 'mcp-sampling'
+
+    @property
+    def system(self) -> str:
+        """The system / model provider, returns `'MCP'`."""
+        return 'MCP'

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pydantic_ai/settings.py

@@ -23,6 +23,7 @@ class ModelSettings(TypedDict, total=False):
     * Cohere
     * Mistral
     * Bedrock
+    * MCP Sampling
     """
 
     temperature: float

{pydantic_ai_slim-0.3.0 → pydantic_ai_slim-0.3.2}/pyproject.toml

@@ -92,7 +92,7 @@ dev = [
     "pytest>=8.3.3",
     "pytest-examples>=0.0.14",
     "pytest-mock>=3.14.0",
-    "pytest-pretty>=1.2.0",
+    "pytest-pretty>=1.3.0",
     "pytest-recording>=0.13.2",
     "diff-cover>=9.2.0",
     "boto3-stubs[bedrock-runtime]",