pydantic-ai-slim 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

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/models/__init__.py

@@ -20,9 +20,12 @@ from typing_extensions import Literal, TypeAliasType, TypedDict
 
 from pydantic_ai.profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
 
+from .. import _utils
+from .._output import OutputObjectDefinition
 from .._parts_manager import ModelResponsePartsManager
 from ..exceptions import UserError
 from ..messages import FileUrl, ModelMessage, ModelRequest, ModelResponse, ModelResponseStreamEvent, VideoUrl
+from ..output import OutputMode
 from ..profiles._json_schema import JsonSchemaTransformer
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -300,13 +303,18 @@ KnownModelName = TypeAliasType(
 """
 
 
-@dataclass
+@dataclass(repr=False)
 class ModelRequestParameters:
     """Configuration for an agent's request to a model, specifically related to tools and output handling."""
 
     function_tools: list[ToolDefinition] = field(default_factory=list)
-    allow_text_output: bool = True
+
+    output_mode: OutputMode = 'text'
+    output_object: OutputObjectDefinition | None = None
     output_tools: list[ToolDefinition] = field(default_factory=list)
+    allow_text_output: bool = True
+
+    __repr__ = _utils.dataclasses_no_defaults_repr
 
 
 class Model(ABC):
@@ -351,6 +359,11 @@ class Model(ABC):
                 function_tools=[_customize_tool_def(transformer, t) for t in model_request_parameters.function_tools],
                 output_tools=[_customize_tool_def(transformer, t) for t in model_request_parameters.output_tools],
             )
+            if output_object := model_request_parameters.output_object:
+                model_request_parameters = replace(
+                    model_request_parameters,
+                    output_object=_customize_output_object(transformer, output_object),
+                )
 
         return model_request_parameters
 
@@ -718,3 +731,9 @@ def _customize_tool_def(transformer: type[JsonSchemaTransformer], t: ToolDefinit
     if t.strict is None:
         t = replace(t, strict=schema_transformer.is_strict_compatible)
     return replace(t, parameters_json_schema=parameters_json_schema)
+
+
+def _customize_output_object(transformer: type[JsonSchemaTransformer], o: OutputObjectDefinition):
+    schema_transformer = transformer(o.json_schema, strict=True)
+    son_schema = schema_transformer.walk()
+    return replace(o, json_schema=son_schema)

pydantic_ai/models/function.py

@@ -11,6 +11,8 @@ from typing import Callable, Union
 
 from typing_extensions import TypeAlias, assert_never, overload
 
+from pydantic_ai.profiles import ModelProfileSpec
+
 from .. import _utils, usage
 from .._utils import PeekableAsyncStream
 from ..messages import (
@@ -49,14 +51,27 @@ class FunctionModel(Model):
     _system: str = field(default='function', repr=False)
 
     @overload
-    def __init__(self, function: FunctionDef, *, model_name: str | None = None) -> None: ...
+    def __init__(
+        self, function: FunctionDef, *, model_name: str | None = None, profile: ModelProfileSpec | None = None
+    ) -> None: ...
 
     @overload
-    def __init__(self, *, stream_function: StreamFunctionDef, model_name: str | None = None) -> None: ...
+    def __init__(
+        self,
+        *,
+        stream_function: StreamFunctionDef,
+        model_name: str | None = None,
+        profile: ModelProfileSpec | None = None,
+    ) -> None: ...
 
     @overload
     def __init__(
-        self, function: FunctionDef, *, stream_function: StreamFunctionDef, model_name: str | None = None
+        self,
+        function: FunctionDef,
+        *,
+        stream_function: StreamFunctionDef,
+        model_name: str | None = None,
+        profile: ModelProfileSpec | None = None,
     ) -> None: ...
 
     def __init__(
@@ -65,6 +80,7 @@ class FunctionModel(Model):
         *,
         stream_function: StreamFunctionDef | None = None,
         model_name: str | None = None,
+        profile: ModelProfileSpec | None = None,
     ):
         """Initialize a `FunctionModel`.
 
@@ -74,6 +90,7 @@ class FunctionModel(Model):
             function: The function to call for non-streamed requests.
             stream_function: The function to call for streamed requests.
             model_name: The name of the model. If not provided, a name is generated from the function names.
+            profile: The model profile to use.
         """
         if function is None and stream_function is None:
             raise TypeError('Either `function` or `stream_function` must be provided')
@@ -83,6 +100,7 @@ class FunctionModel(Model):
         function_name = self.function.__name__ if self.function is not None else ''
         stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
         self._model_name = model_name or f'function:{function_name}:{stream_function_name}'
+        self._profile = profile
 
     async def request(
         self,

pydantic_ai/models/gemini.py

@@ -16,6 +16,8 @@ from typing_extensions import NotRequired, TypedDict, assert_never
 from pydantic_ai.providers import Provider, infer_provider
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
+from .._output import OutputObjectDefinition
+from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
     FileUrl,
@@ -203,12 +205,10 @@ class GeminiModel(Model):
     def _get_tool_config(
         self, model_request_parameters: ModelRequestParameters, tools: _GeminiTools | None
     ) -> _GeminiToolConfig | None:
-        if model_request_parameters.allow_text_output:
-            return None
-        elif tools:
+        if not model_request_parameters.allow_text_output and tools:
             return _tool_config([t['name'] for t in tools['function_declarations']])
         else:
-            return _tool_config([])  # pragma: no cover
+            return None
 
     @asynccontextmanager
     async def _make_request(
@@ -231,6 +231,18 @@ class GeminiModel(Model):
             request_data['toolConfig'] = tool_config
 
         generation_config = _settings_to_generation_config(model_settings)
+        if model_request_parameters.output_mode == 'native':
+            if tools:
+                raise UserError('Gemini does not support structured output and tools at the same time.')
+
+            generation_config['response_mime_type'] = 'application/json'
+
+            output_object = model_request_parameters.output_object
+            assert output_object is not None
+            generation_config['response_schema'] = self._map_response_schema(output_object)
+        elif model_request_parameters.output_mode == 'prompted' and not tools:
+            generation_config['response_mime_type'] = 'application/json'
+
         if generation_config:
             request_data['generationConfig'] = generation_config
 
@@ -376,6 +388,15 @@ class GeminiModel(Model):
                     assert_never(item)
         return content
 
+    def _map_response_schema(self, o: OutputObjectDefinition) -> dict[str, Any]:
+        response_schema = o.json_schema.copy()
+        if o.name:
+            response_schema['title'] = o.name
+        if o.description:
+            response_schema['description'] = o.description
+
+        return response_schema
+
 
 def _settings_to_generation_config(model_settings: GeminiModelSettings) -> _GeminiGenerationConfig:
     config: _GeminiGenerationConfig = {}
@@ -577,6 +598,8 @@ class _GeminiGenerationConfig(TypedDict, total=False):
     frequency_penalty: float
     stop_sequences: list[str]
     thinking_config: ThinkingConfig
+    response_mime_type: str
+    response_schema: dict[str, Any]
 
 
 class _GeminiContent(TypedDict):