pydantic-ai-slim 0.0.40__tar.gz → 0.0.42__tar.gz

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/.gitignore

@@ -1,5 +1,4 @@
 site
-.python-version
 .venv
 dist
 __pycache__
@@ -16,3 +15,4 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 .vscode/
 /question_graph_history.json
 /docs-site/.wrangler/
+/CLAUDE.md

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.40
+Version: 0.0.42
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,7 +29,7 @@ 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.0.40
+Requires-Dist: pydantic-graph==0.0.42
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
@@ -45,9 +45,11 @@ Requires-Dist: cohere>=5.13.11; extra == 'cohere'
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
 Provides-Extra: groq
-Requires-Dist: groq>=0.12.0; extra == 'groq'
+Requires-Dist: groq>=0.15.0; extra == 'groq'
 Provides-Extra: logfire
 Requires-Dist: logfire>=2.3; extra == 'logfire'
+Provides-Extra: mcp
+Requires-Dist: mcp>=1.4.1; (python_version >= '3.10') and extra == 'mcp'
 Provides-Extra: mistral
 Requires-Dist: mistralai>=1.2.5; extra == 'mistral'
 Provides-Extra: openai

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/_agent_graph.py

@@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import field
-from typing import Any, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast
 
 from opentelemetry.trace import Span, Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
@@ -27,11 +27,10 @@ from . import (
 from .models.instrumented import InstrumentedModel
 from .result import ResultDataT
 from .settings import ModelSettings, merge_model_settings
-from .tools import (
-    RunContext,
-    Tool,
-    ToolDefinition,
-)
+from .tools import RunContext, Tool, ToolDefinition
+
+if TYPE_CHECKING:
+    from .mcp import MCPServer
 
 __all__ = (
     'GraphAgentState',
@@ -94,6 +93,7 @@ class GraphAgentDeps(Generic[DepsT, ResultDataT]):
     result_validators: list[_result.ResultValidator[DepsT, ResultDataT]]
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
+    mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
 
     run_span: Span
     tracer: Tracer
@@ -219,7 +219,17 @@ async def _prepare_request_parameters(
         if tool_def := await tool.prepare_tool_def(ctx):
             function_tool_defs.append(tool_def)
 
-    await asyncio.gather(*map(add_tool, ctx.deps.function_tools.values()))
+    async def add_mcp_server_tools(server: MCPServer) -> None:
+        if not server.is_running:
+            raise exceptions.UserError(f'MCP server is not running: {server}')
+        tool_defs = await server.list_tools()
+        # TODO(Marcelo): We should check if the tool names are unique. If not, we should raise an error.
+        function_tool_defs.extend(tool_defs)
+
+    await asyncio.gather(
+        *map(add_tool, ctx.deps.function_tools.values()),
+        *map(add_mcp_server_tools, ctx.deps.mcp_servers),
+    )
 
     result_schema = ctx.deps.result_schema
     return models.ModelRequestParameters(
@@ -594,6 +604,21 @@ async def process_function_tools(
                 yield event
                 call_index_to_event_id[len(calls_to_run)] = event.call_id
                 calls_to_run.append((tool, call))
+        elif mcp_tool := await _tool_from_mcp_server(call.tool_name, ctx):
+            if stub_function_tools:
+                # TODO(Marcelo): We should add coverage for this part of the code.
+                output_parts.append(  # pragma: no cover
+                    _messages.ToolReturnPart(
+                        tool_name=call.tool_name,
+                        content='Tool not executed - a final result was already processed.',
+                        tool_call_id=call.tool_call_id,
+                    )
+                )
+            else:
+                event = _messages.FunctionToolCallEvent(call)
+                yield event
+                call_index_to_event_id[len(calls_to_run)] = event.call_id
+                calls_to_run.append((mcp_tool, call))
         elif result_schema is not None and call.tool_name in result_schema.tools:
             # if tool_name is in _result_schema, it means we found a result tool but an error occurred in
             # validation, we don't add another part here
@@ -641,6 +666,35 @@ async def process_function_tools(
         output_parts.append(results_by_index[k])
 
 
+async def _tool_from_mcp_server(
+    tool_name: str,
+    ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+) -> Tool[DepsT] | None:
+    """Call each MCP server to find the tool with the given name.
+
+    Args:
+        tool_name: The name of the tool to find.
+        ctx: The current run context.
+
+    Returns:
+        The tool with the given name, or `None` if no tool with the given name is found.
+    """
+
+    async def run_tool(ctx: RunContext[DepsT], **args: Any) -> Any:
+        # There's no normal situation where the server will not be running at this point, we check just in case
+        # some weird edge case occurs.
+        if not server.is_running:  # pragma: no cover
+            raise exceptions.UserError(f'MCP server is not running: {server}')
+        result = await server.call_tool(tool_name, args)
+        return result
+
+    for server in ctx.deps.mcp_servers:
+        tools = await server.list_tools()
+        if tool_name in {tool.name for tool in tools}:
+            return Tool(name=tool_name, function=run_tool, takes_ctx=True)
+    return None
+
+
 def _unknown_tool(
     tool_name: str,
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/_pydantic.py

@@ -44,6 +44,7 @@ def function_schema(  # noqa: C901
     takes_ctx: bool,
     docstring_format: DocstringFormat,
     require_parameter_descriptions: bool,
+    schema_generator: type[GenerateJsonSchema],
 ) -> FunctionSchema:
     """Build a Pydantic validator and JSON schema from a tool function.
 
@@ -52,6 +53,7 @@ def function_schema(  # noqa: C901
         takes_ctx: Whether the function takes a `RunContext` first argument.
         docstring_format: The docstring format to use.
         require_parameter_descriptions: Whether to require descriptions for all tool function parameters.
+        schema_generator: The JSON schema generator class to use.
 
     Returns:
         A `FunctionSchema` instance.
@@ -150,14 +152,12 @@ def function_schema(  # noqa: C901
     )
     # PluggableSchemaValidator is api compatible with SchemaValidator
     schema_validator = cast(SchemaValidator, schema_validator)
-    json_schema = GenerateJsonSchema().generate(schema)
+    json_schema = schema_generator().generate(schema)
 
     # workaround for https://github.com/pydantic/pydantic/issues/10785
-    # if we build a custom TypeDict schema (matches when `single_arg_name is None`), we manually set
+    # if we build a custom TypedDict schema (matches when `single_arg_name is None`), we manually set
     # `additionalProperties` in the JSON Schema
-    if single_arg_name is None:
-        json_schema['additionalProperties'] = bool(var_kwargs_schema)
-    elif not description:
+    if single_arg_name is not None and not description:
         # if the tool description is not set, and we have a single parameter, take the description from that
         # and set it on the tool
         description = json_schema.pop('description', None)
@@ -218,6 +218,7 @@ def _build_schema(
     td_schema = core_schema.typed_dict_schema(
         fields,
         config=core_config,
+        total=var_kwargs_schema is None,
         extras_schema=gen_schema.generate_schema(var_kwargs_schema) if var_kwargs_schema else None,
     )
     return td_schema, None

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/agent.py

@@ -3,12 +3,13 @@ from __future__ import annotations as _annotations
 import dataclasses
 import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
-from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
 from copy import deepcopy
 from types import FrameType
-from typing import Any, Callable, ClassVar, Generic, cast, final, overload
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, cast, final, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
+from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import TypeGuard, TypeVar, deprecated
 
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
@@ -31,6 +32,7 @@ from .settings import ModelSettings, merge_model_settings
 from .tools import (
     AgentDepsT,
     DocstringFormat,
+    GenerateToolJsonSchema,
     RunContext,
     Tool,
     ToolFuncContext,
@@ -47,6 +49,9 @@ CallToolsNode = _agent_graph.CallToolsNode
 ModelRequestNode = _agent_graph.ModelRequestNode
 UserPromptNode = _agent_graph.UserPromptNode
 
+if TYPE_CHECKING:
+    from pydantic_ai.mcp import MCPServer
+
 __all__ = (
     'Agent',
     'AgentRun',
@@ -129,6 +134,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         repr=False
     )
     _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
+    _mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
     _default_retries: int = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
     _override_deps: _utils.Option[AgentDepsT] = dataclasses.field(default=None, repr=False)
@@ -148,6 +154,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         result_tool_description: str | None = None,
         result_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
@@ -173,6 +180,8 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             result_retries: The maximum number of retries to allow for result validation, defaults to `retries`.
             tools: Tools to register with the agent, you can also register tools via the decorators
                 [`@agent.tool`][pydantic_ai.Agent.tool] and [`@agent.tool_plain`][pydantic_ai.Agent.tool_plain].
+            mcp_servers: MCP servers to register with the agent. You should register a [`MCPServer`][pydantic_ai.mcp.MCPServer]
+                for each server you want the agent to connect to.
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
                 it's evaluated to create a [`Model`][pydantic_ai.models.Model] instance immediately,
                 which checks for the necessary environment variables. Set this to `false`
@@ -215,6 +224,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         self._default_retries = retries
         self._max_result_retries = result_retries if result_retries is not None else retries
+        self._mcp_servers = mcp_servers
         for tool in tools:
             if isinstance(tool, Tool):
                 self._register_tool(tool)
@@ -461,6 +471,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             result_tools=self._result_schema.tool_defs() if self._result_schema else [],
             result_validators=result_validators,
             function_tools=self._function_tools,
+            mcp_servers=self._mcp_servers,
             run_span=run_span,
             tracer=tracer,
         )
@@ -927,6 +938,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ...
 
     def tool(
@@ -939,6 +951,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Any:
         """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.
 
@@ -980,6 +993,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
                 Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
             require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
         """
         if func is None:
 
@@ -988,7 +1002,14 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             ) -> ToolFuncContext[AgentDepsT, ToolParams]:
                 # noinspection PyTypeChecker
                 self._register_function(
-                    func_, True, name, retries, prepare, docstring_format, require_parameter_descriptions
+                    func_,
+                    True,
+                    name,
+                    retries,
+                    prepare,
+                    docstring_format,
+                    require_parameter_descriptions,
+                    schema_generator,
                 )
                 return func_
 
@@ -996,7 +1017,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             # noinspection PyTypeChecker
             self._register_function(
-                func, True, name, retries, prepare, docstring_format, require_parameter_descriptions
+                func, True, name, retries, prepare, docstring_format, require_parameter_descriptions, schema_generator
             )
             return func
 
@@ -1013,6 +1034,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ...
 
     def tool_plain(
@@ -1025,6 +1047,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
@@ -1066,20 +1089,28 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
                 Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
             require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+            schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
         """
         if func is None:
 
             def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
                 # noinspection PyTypeChecker
                 self._register_function(
-                    func_, False, name, retries, prepare, docstring_format, require_parameter_descriptions
+                    func_,
+                    False,
+                    name,
+                    retries,
+                    prepare,
+                    docstring_format,
+                    require_parameter_descriptions,
+                    schema_generator,
                 )
                 return func_
 
             return tool_decorator
         else:
             self._register_function(
-                func, False, name, retries, prepare, docstring_format, require_parameter_descriptions
+                func, False, name, retries, prepare, docstring_format, require_parameter_descriptions, schema_generator
             )
             return func
 
@@ -1092,6 +1123,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         prepare: ToolPrepareFunc[AgentDepsT] | None,
         docstring_format: DocstringFormat,
         require_parameter_descriptions: bool,
+        schema_generator: type[GenerateJsonSchema],
     ) -> None:
         """Private utility to register a function as a tool."""
         retries_ = retries if retries is not None else self._default_retries
@@ -1103,6 +1135,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             prepare=prepare,
             docstring_format=docstring_format,
             require_parameter_descriptions=require_parameter_descriptions,
+            schema_generator=schema_generator,
         )
         self._register_tool(tool)
 
@@ -1253,6 +1286,20 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         """
         return isinstance(node, End)
 
+    @asynccontextmanager
+    async def run_mcp_servers(self) -> 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.
+        """
+        exit_stack = AsyncExitStack()
+        try:
+            for mcp_server in self._mcp_servers:
+                await exit_stack.enter_async_context(mcp_server)
+            yield
+        finally:
+            await exit_stack.aclose()
+
 
 @dataclasses.dataclass(repr=False)
 class AgentRun(Generic[AgentDepsT, ResultDataT]):

pydantic_ai_slim-0.0.42/pydantic_ai/mcp.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Sequence
+from contextlib import AsyncExitStack, asynccontextmanager
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any
+
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.types import JSONRPCMessage
+from typing_extensions import Self
+
+from pydantic_ai.tools import ToolDefinition
+
+try:
+    from mcp.client.session import ClientSession
+    from mcp.client.sse import sse_client
+    from mcp.client.stdio import StdioServerParameters, stdio_client
+    from mcp.types import CallToolResult
+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
+
+__all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP'
+
+
+class MCPServer(ABC):
+    """Base class for attaching agents to MCP servers.
+
+    See <https://modelcontextprotocol.io> for more information.
+    """
+
+    is_running: bool = False
+
+    _client: ClientSession
+    _read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
+    _write_stream: MemoryObjectSendStream[JSONRPCMessage]
+    _exit_stack: AsyncExitStack
+
+    @abstractmethod
+    @asynccontextmanager
+    async def client_streams(
+        self,
+    ) -> AsyncIterator[
+        tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]
+    ]:
+        """Create the streams for the MCP server."""
+        raise NotImplementedError('MCP Server subclasses must implement this method.')
+        yield
+
+    async def list_tools(self) -> list[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()
+        return [
+            ToolDefinition(
+                name=tool.name,
+                description=tool.description or '',
+                parameters_json_schema=tool.inputSchema,
+            )
+            for tool in tools.tools
+        ]
+
+    async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> CallToolResult:
+        """Call a tool on the server.
+
+        Args:
+            tool_name: The name of the tool to call.
+            arguments: The arguments to pass to the tool.
+
+        Returns:
+            The result of the tool call.
+        """
+        return await self._client.call_tool(tool_name, arguments)
+
+    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)
+
+        await self._client.initialize()
+        self.is_running = True
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
+    ) -> bool | None:
+        await self._exit_stack.aclose()
+        self.is_running = False
+
+
+@dataclass
+class MCPServerStdio(MCPServer):
+    """Runs an MCP server in a subprocess and communicates with it over stdin/stdout.
+
+    This class implements the stdio transport from the MCP specification.
+    See <https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio> for more information.
+
+    !!! note
+        Using this class as an async context manager will start the server as a subprocess when entering the context,
+        and stop it when exiting the context.
+
+    Example:
+    ```python {py="3.10"}
+    from pydantic_ai import Agent
+    from pydantic_ai.mcp import MCPServerStdio
+
+    server = MCPServerStdio('npx', ['-y', '@pydantic/mcp-run-python', 'stdio'])  # (1)!
+    agent = Agent('openai:gpt-4o', mcp_servers=[server])
+
+    async def main():
+        async with agent.run_mcp_servers():  # (2)!
+            ...
+    ```
+
+    1. See [MCP Run Python](../mcp/run-python.md) for more information.
+    2. This will start the server as a subprocess and connect to it.
+    """
+
+    command: str
+    """The command to run."""
+
+    args: Sequence[str]
+    """The arguments to pass to the command."""
+
+    env: dict[str, str] | None = 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`.
+    """
+
+    @asynccontextmanager
+    async def client_streams(
+        self,
+    ) -> AsyncIterator[
+        tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]
+    ]:
+        server = StdioServerParameters(command=self.command, args=list(self.args), env=self.env)
+        async with stdio_client(server=server) as (read_stream, write_stream):
+            yield read_stream, write_stream
+
+
+@dataclass
+class MCPServerHTTP(MCPServer):
+    """An MCP server that connects over streamable HTTP connections.
+
+    This class implements the SSE transport from the MCP specification.
+    See <https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse> for more information.
+
+    The name "HTTP" is used since this implemented will be adapted in future to use the new
+    [Streamable HTTP](https://github.com/modelcontextprotocol/specification/pull/206) currently in development.
+
+    !!! note
+        Using this class as an async context manager will create a new pool of HTTP connections to connect
+        to a server which should already be running.
+
+    Example:
+    ```python {py="3.10"}
+    from pydantic_ai import Agent
+    from pydantic_ai.mcp import MCPServerHTTP
+
+    server = MCPServerHTTP('http://localhost:3001/sse')  # (1)!
+    agent = Agent('openai:gpt-4o', mcp_servers=[server])
+
+    async def main():
+        async with agent.run_mcp_servers():  # (2)!
+            ...
+    ```
+
+    1. E.g. you might be connecting to a server run with `npx @pydantic/mcp-run-python sse`,
+      see [MCP Run Python](../mcp/run-python.md) for more information.
+    2. This will connect to a server running on `localhost:3001`.
+    """
+
+    url: str
+    """The URL of the SSE endpoint on the MCP server.
+
+    For example for a server running locally, this might be `http://localhost:3001/sse`.
+    """
+
+    @asynccontextmanager
+    async def client_streams(
+        self,
+    ) -> AsyncIterator[
+        tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]
+    ]:  # pragma: no cover
+        async with sse_client(url=self.url) as (read_stream, write_stream):
+            yield read_stream, write_stream

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/models/anthropic.py

@@ -11,7 +11,7 @@ from typing import Any, Literal, Union, cast, overload
 
 from anthropic.types import DocumentBlockParam
 from httpx import AsyncClient as AsyncHTTPClient
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
 from .._utils import guard_tool_call_id as _guard_tool_call_id
@@ -31,6 +31,7 @@ from ..messages import (
     ToolReturnPart,
     UserPromptPart,
 )
+from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, cached_async_http_client, check_allow_model_requests
@@ -111,10 +112,31 @@ class AnthropicModel(Model):
     _model_name: AnthropicModelName = field(repr=False)
     _system: str = field(default='anthropic', repr=False)
 
+    @overload
+    def __init__(
+        self,
+        model_name: AnthropicModelName,
+        *,
+        provider: Literal['anthropic'] | Provider[AsyncAnthropic] = 'anthropic',
+    ) -> None: ...
+
+    @deprecated('Use the `provider` parameter instead of `api_key`, `anthropic_client`, and `http_client`.')
+    @overload
+    def __init__(
+        self,
+        model_name: AnthropicModelName,
+        *,
+        provider: None = None,
+        api_key: str | None = None,
+        anthropic_client: AsyncAnthropic | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None: ...
+
     def __init__(
         self,
         model_name: AnthropicModelName,
         *,
+        provider: Literal['anthropic'] | Provider[AsyncAnthropic] | None = None,
         api_key: str | None = None,
         anthropic_client: AsyncAnthropic | None = None,
         http_client: AsyncHTTPClient | None = None,
@@ -124,6 +146,8 @@ class AnthropicModel(Model):
         Args:
             model_name: The name of the Anthropic model to use. List of model names available
                 [here](https://docs.anthropic.com/en/docs/about-claude/models).
+            provider: The provider to use for the Anthropic API. Can be either the string 'anthropic' or an
+                instance of `Provider[AsyncAnthropic]`. If not provided, the other parameters will be used.
             api_key: The API key to use for authentication, if not provided, the `ANTHROPIC_API_KEY` environment variable
                 will be used if available.
             anthropic_client: An existing
@@ -132,7 +156,12 @@ class AnthropicModel(Model):
             http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
         """
         self._model_name = model_name
-        if anthropic_client is not None:
+
+        if provider is not None:
+            if isinstance(provider, str):
+                provider = infer_provider(provider)
+            self.client = provider.client
+        elif anthropic_client is not None:
             assert http_client is None, 'Cannot provide both `anthropic_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `anthropic_client` and `api_key`'
             self.client = anthropic_client

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/models/fallback.py

@@ -70,14 +70,9 @@ class FallbackModel(Model):
                     exceptions.append(exc)
                     continue
                 raise exc
-            else:
-                with suppress(Exception):
-                    span = get_current_span()
-                    if span.is_recording():
-                        attributes = getattr(span, 'attributes', {})
-                        if attributes.get('gen_ai.request.model') == self.model_name:
-                            span.set_attributes(InstrumentedModel.model_attributes(model))
-                return response, usage
+
+            self._set_span_attributes(model)
+            return response, usage
 
         raise FallbackExceptionGroup('All models from FallbackModel failed', exceptions)
 
@@ -102,11 +97,21 @@ class FallbackModel(Model):
                         exceptions.append(exc)
                         continue
                     raise exc
+
+                self._set_span_attributes(model)
                 yield response
                 return
 
         raise FallbackExceptionGroup('All models from FallbackModel failed', exceptions)
 
+    def _set_span_attributes(self, model: Model):
+        with suppress(Exception):
+            span = get_current_span()
+            if span.is_recording():
+                attributes = getattr(span, 'attributes', {})
+                if attributes.get('gen_ai.request.model') == self.model_name:
+                    span.set_attributes(InstrumentedModel.model_attributes(model))
+
     @property
     def model_name(self) -> str:
         """The model name."""

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/models/gemini.py

@@ -139,11 +139,9 @@ class GeminiModel(Model):
 
         if provider is not None:
             if isinstance(provider, str):
-                self._system = provider
-                self.client = infer_provider(provider).client
-            else:
-                self._system = provider.name
-                self.client = provider.client
+                provider = infer_provider(provider)
+            self._system = provider.name
+            self.client = provider.client
             self._url = str(self.client.base_url)
         else:
             if api_key is None:

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/models/groq.py

@@ -138,9 +138,8 @@ class GroqModel(Model):
 
         if provider is not None:
             if isinstance(provider, str):
-                self.client = infer_provider(provider).client
-            else:
-                self.client = provider.client
+                provider = infer_provider(provider)
+            self.client = provider.client
         elif groq_client is not None:
             assert http_client is None, 'Cannot provide both `groq_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `groq_client` and `api_key`'

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/models/mistral.py

@@ -7,11 +7,11 @@ from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from itertools import chain
-from typing import Any, Callable, Literal, Union, cast
+from typing import Any, Callable, Literal, Union, cast, overload
 
 import pydantic_core
 from httpx import AsyncClient as AsyncHTTPClient, Timeout
-from typing_extensions import assert_never
+from typing_extensions import assert_never, deprecated
 
 from .. import ModelHTTPError, UnexpectedModelBehavior, _utils
 from .._utils import now_utc as _now_utc
@@ -31,6 +31,7 @@ from ..messages import (
     ToolReturnPart,
     UserPromptPart,
 )
+from ..providers import Provider, infer_provider
 from ..result import Usage
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
@@ -112,10 +113,33 @@ class MistralModel(Model):
     _model_name: MistralModelName = field(repr=False)
     _system: str = field(default='mistral_ai', repr=False)
 
+    @overload
     def __init__(
         self,
         model_name: MistralModelName,
         *,
+        provider: Literal['mistral'] | Provider[Mistral] = 'mistral',
+        json_mode_schema_prompt: str = """Answer in JSON Object, respect the format:\n```\n{schema}\n```\n""",
+    ) -> None: ...
+
+    @overload
+    @deprecated('Use the `provider` parameter instead of `api_key`, `client` and `http_client`.')
+    def __init__(
+        self,
+        model_name: MistralModelName,
+        *,
+        provider: None = None,
+        api_key: str | Callable[[], str | None] | None = None,
+        client: Mistral | None = None,
+        http_client: AsyncHTTPClient | None = None,
+        json_mode_schema_prompt: str = """Answer in JSON Object, respect the format:\n```\n{schema}\n```\n""",
+    ) -> None: ...
+
+    def __init__(
+        self,
+        model_name: MistralModelName,
+        *,
+        provider: Literal['mistral'] | Provider[Mistral] | None = None,
         api_key: str | Callable[[], str | None] | None = None,
         client: Mistral | None = None,
         http_client: AsyncHTTPClient | None = None,
@@ -124,6 +148,9 @@ class MistralModel(Model):
         """Initialize a Mistral model.
 
         Args:
+            provider: The provider to use for authentication and API access. Can be either the string
+                'mistral' or an instance of `Provider[Mistral]`. If not provided, a new provider will be
+                created using the other parameters.
             model_name: The name of the model to use.
             api_key: The API key to use for authentication, if unset uses `MISTRAL_API_KEY` environment variable.
             client: An existing `Mistral` client to use, if provided, `api_key` and `http_client` must be `None`.
@@ -133,17 +160,22 @@ class MistralModel(Model):
         self._model_name = model_name
         self.json_mode_schema_prompt = json_mode_schema_prompt
 
-        if client is not None:
+        if provider is not None:
+            if isinstance(provider, str):
+                # TODO(Marcelo): We should add an integration test with VCR when I get the API key.
+                provider = infer_provider(provider)  # pragma: no cover
+            self.client = provider.client
+        elif client is not None:
             assert http_client is None, 'Cannot provide both `mistral_client` and `http_client`'
             assert api_key is None, 'Cannot provide both `mistral_client` and `api_key`'
             self.client = client
         else:
-            api_key = os.getenv('MISTRAL_API_KEY') if api_key is None else api_key
+            api_key = api_key or os.getenv('MISTRAL_API_KEY')
             self.client = Mistral(api_key=api_key, async_client=http_client or cached_async_http_client())
 
     @property
     def base_url(self) -> str:
-        return str(self.client.sdk_configuration.get_server_details()[0])
+        return self.client.sdk_configuration.get_server_details()[0]
 
     async def request(
         self,

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/models/openai.py

@@ -162,9 +162,8 @@ class OpenAIModel(Model):
 
         if provider is not None:
             if isinstance(provider, str):
-                self.client = infer_provider(provider).client
-            else:
-                self.client = provider.client
+                provider = infer_provider(provider)
+            self.client = provider.client
         else:  # pragma: no cover
             # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
             # openai compatible models do not always need an API key, but a placeholder (non-empty) key is required.

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/providers/__init__.py

@@ -69,5 +69,13 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .groq import GroqProvider
 
         return GroqProvider()
+    elif provider == 'anthropic':
+        from .anthropic import AnthropicProvider
+
+        return AnthropicProvider()
+    elif provider == 'mistral':
+        from .mistral import MistralProvider
+
+        return MistralProvider()
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')

pydantic_ai_slim-0.0.42/pydantic_ai/providers/anthropic.py

@@ -0,0 +1,74 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+import httpx
+
+from pydantic_ai.models import cached_async_http_client
+
+try:
+    from anthropic import AsyncAnthropic
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `anthropic` package to use the Anthropic provider, '
+        "you can use the `anthropic` optional group — `pip install 'pydantic-ai-slim[anthropic]'`"
+    ) from _import_error
+
+
+from . import Provider
+
+
+class AnthropicProvider(Provider[AsyncAnthropic]):
+    """Provider for Anthropic API."""
+
+    @property
+    def name(self) -> str:
+        return 'anthropic'
+
+    @property
+    def base_url(self) -> str:
+        return str(self._client.base_url)
+
+    @property
+    def client(self) -> AsyncAnthropic:
+        return self._client
+
+    @overload
+    def __init__(self, *, anthropic_client: AsyncAnthropic | None = None) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str | None = None, http_client: httpx.AsyncClient | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        anthropic_client: AsyncAnthropic | None = None,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        """Create a new Anthropic provider.
+
+        Args:
+            api_key: The API key to use for authentication, if not provided, the `ANTHROPIC_API_KEY` environment variable
+                will be used if available.
+            anthropic_client: An existing [`AsyncAnthropic`](https://github.com/anthropics/anthropic-sdk-python)
+                client to use. If provided, the `api_key` and `http_client` arguments will be ignored.
+            http_client: An existing `httpx.AsyncClient` to use for making HTTP requests.
+        """
+        if anthropic_client is not None:
+            assert http_client is None, 'Cannot provide both `anthropic_client` and `http_client`'
+            assert api_key is None, 'Cannot provide both `anthropic_client` and `api_key`'
+            self._client = anthropic_client
+        else:
+            api_key = api_key or os.environ.get('ANTHROPIC_API_KEY')
+            if api_key is None:
+                raise ValueError(
+                    'Set the `ANTHROPIC_API_KEY` environment variable or pass it via `AnthropicProvider(api_key=...)`'
+                    'to use the Anthropic provider.'
+                )
+
+            if http_client is not None:
+                self._client = AsyncAnthropic(api_key=api_key, http_client=http_client)
+            else:
+                self._client = AsyncAnthropic(api_key=api_key, http_client=cached_async_http_client())

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/providers/groq.py

@@ -57,17 +57,21 @@ class GroqProvider(Provider[AsyncGroq]):
                 client to use. If provided, `api_key` and `http_client` must be `None`.
             http_client: An existing `AsyncHTTPClient` to use for making HTTP requests.
         """
-        api_key = api_key or os.environ.get('GROQ_API_KEY')
-
-        if api_key is None and groq_client is None:
-            raise ValueError(
-                'Set the `GROQ_API_KEY` environment variable or pass it via `GroqProvider(api_key=...)`'
-                'to use the Groq provider.'
-            )
-
         if groq_client is not None:
+            assert http_client is None, 'Cannot provide both `groq_client` and `http_client`'
+            assert api_key is None, 'Cannot provide both `groq_client` and `api_key`'
             self._client = groq_client
-        elif http_client is not None:
-            self._client = AsyncGroq(base_url=self.base_url, api_key=api_key, http_client=http_client)
         else:
-            self._client = AsyncGroq(base_url=self.base_url, api_key=api_key, http_client=cached_async_http_client())
+            api_key = api_key or os.environ.get('GROQ_API_KEY')
+
+            if api_key is None:
+                raise ValueError(
+                    'Set the `GROQ_API_KEY` environment variable or pass it via `GroqProvider(api_key=...)`'
+                    'to use the Groq provider.'
+                )
+            elif http_client is not None:
+                self._client = AsyncGroq(base_url=self.base_url, api_key=api_key, http_client=http_client)
+            else:
+                self._client = AsyncGroq(
+                    base_url=self.base_url, api_key=api_key, http_client=cached_async_http_client()
+                )

pydantic_ai_slim-0.0.42/pydantic_ai/providers/mistral.py

@@ -0,0 +1,73 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+
+from pydantic_ai.models import cached_async_http_client
+
+try:
+    from mistralai import Mistral
+except ImportError as e:  # pragma: no cover
+    raise ImportError(
+        'Please install the `mistral` package to use the Mistral provider, '
+        "you can use the `mistral` optional group — `pip install 'pydantic-ai-slim[mistral]'`"
+    ) from e
+
+
+from . import Provider
+
+
+class MistralProvider(Provider[Mistral]):
+    """Provider for Mistral API."""
+
+    @property
+    def name(self) -> str:
+        return 'mistral'
+
+    @property
+    def base_url(self) -> str:
+        return self.client.sdk_configuration.get_server_details()[0]
+
+    @property
+    def client(self) -> Mistral:
+        return self._client
+
+    @overload
+    def __init__(self, *, mistral_client: Mistral | None = None) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str | None = None, http_client: AsyncHTTPClient | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        mistral_client: Mistral | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        """Create a new Mistral provider.
+
+        Args:
+            api_key: The API key to use for authentication, if not provided, the `MISTRAL_API_KEY` environment variable
+                will be used if available.
+            mistral_client: An existing `Mistral` client to use, if provided, `api_key` and `http_client` must be `None`.
+            http_client: An existing async client to use for making HTTP requests.
+        """
+        if mistral_client is not None:
+            assert http_client is None, 'Cannot provide both `mistral_client` and `http_client`'
+            assert api_key is None, 'Cannot provide both `mistral_client` and `api_key`'
+            self._client = mistral_client
+        else:
+            api_key = api_key or os.environ.get('MISTRAL_API_KEY')
+
+            if api_key is None:
+                raise ValueError(
+                    'Set the `MISTRAL_API_KEY` environment variable or pass it via `MistralProvider(api_key=...)`'
+                    'to use the Mistral provider.'
+                )
+            elif http_client is not None:
+                self._client = Mistral(api_key=api_key, async_client=http_client)
+            else:
+                self._client = Mistral(api_key=api_key, async_client=cached_async_http_client())

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pydantic_ai/tools.py

@@ -7,7 +7,8 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
 
 from pydantic import ValidationError
-from pydantic_core import SchemaValidator
+from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
+from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import Concatenate, ParamSpec, TypeAlias, TypeVar
 
 from . import _pydantic, _utils, messages as _messages, models
@@ -142,6 +143,22 @@ DocstringFormat = Literal['google', 'numpy', 'sphinx', 'auto']
 A = TypeVar('A')
 
 
+class GenerateToolJsonSchema(GenerateJsonSchema):
+    def typed_dict_schema(self, schema: core_schema.TypedDictSchema) -> JsonSchemaValue:
+        s = super().typed_dict_schema(schema)
+        total = schema.get('total')
+        if total is not None:
+            s['additionalProperties'] = not total
+        return s
+
+    def _named_required_fields_schema(self, named_required_fields: Sequence[tuple[str, bool, Any]]) -> JsonSchemaValue:
+        # Remove largely-useless property titles
+        s = super()._named_required_fields_schema(named_required_fields)
+        for p in s.get('properties', {}):
+            s['properties'][p].pop('title', None)
+        return s
+
+
 @dataclass(init=False)
 class Tool(Generic[AgentDepsT]):
     """A tool function for an agent."""
@@ -176,6 +193,7 @@ class Tool(Generic[AgentDepsT]):
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
+        schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
     ):
         """Create a new tool instance.
 
@@ -225,11 +243,14 @@ class Tool(Generic[AgentDepsT]):
             docstring_format: The format of the docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
                 Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
             require_parameter_descriptions: If True, raise an error if a parameter description is missing. Defaults to False.
+            schema_generator: The JSON schema generator class to use. Defaults to `GenerateToolJsonSchema`.
         """
         if takes_ctx is None:
             takes_ctx = _pydantic.takes_ctx(function)
 
-        f = _pydantic.function_schema(function, takes_ctx, docstring_format, require_parameter_descriptions)
+        f = _pydantic.function_schema(
+            function, takes_ctx, docstring_format, require_parameter_descriptions, schema_generator
+        )
         self.function = function
         self.takes_ctx = takes_ctx
         self.max_retries = max_retries

{pydantic_ai_slim-0.0.40 → pydantic_ai_slim-0.0.42}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.40"
+version = "0.0.42"
 description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
 authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
 license = "MIT"
@@ -36,7 +36,7 @@ dependencies = [
     "griffe>=1.3.2",
     "httpx>=0.27",
     "pydantic>=2.10",
-    "pydantic-graph==0.0.40",
+    "pydantic-graph==0.0.42",
     "exceptiongroup; python_version < '3.11'",
     "opentelemetry-api>=1.28.0",
     "typing-inspection>=0.4.0",
@@ -50,7 +50,7 @@ openai = ["openai>=1.65.1"]
 cohere = ["cohere>=5.13.11"]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.3"]
 anthropic = ["anthropic>=0.49.0"]
-groq = ["groq>=0.12.0"]
+groq = ["groq>=0.15.0"]
 mistral = ["mistralai>=1.2.5"]
 bedrock = ["boto3>=1.34.116"]
 # Tools
@@ -58,6 +58,8 @@ duckduckgo = ["duckduckgo-search>=7.0.0"]
 tavily = ["tavily-python>=0.5.0"]
 # CLI
 cli = ["rich>=13", "prompt-toolkit>=3", "argcomplete>=3.5.0"]
+# MCP
+mcp = ["mcp>=1.4.1; python_version >= '3.10'"]
 
 [dependency-groups]
 dev = [
@@ -75,6 +77,9 @@ dev = [
     "boto3-stubs[bedrock-runtime]",
 ]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [project.scripts]
 pai = "pydantic_ai._cli:app"