pydantic-ai-slim 0.0.41__tar.gz → 0.0.43__tar.gz

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/.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.41 → pydantic_ai_slim-0.0.43}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.41
+Version: 0.0.43
 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.41
+Requires-Dist: pydantic-graph==0.0.43
 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.41 → pydantic_ai_slim-0.0.43}/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.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/_cli.py

@@ -31,7 +31,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `rich`, `prompt-toolkit` and `argcomplete` to use the PydanticAI CLI, '
-        "you can use the `cli` optional group — `pip install 'pydantic-ai-slim[cli]'`"
+        'you can use the `cli` optional group — `pip install "pydantic-ai-slim[cli]"`'
     ) from _import_error
 
 from pydantic_ai.agent import Agent

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/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.41 → pydantic_ai_slim-0.0.43}/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`
@@ -186,6 +195,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                 If this isn't set, then the last value set by
                 [`Agent.instrument_all()`][pydantic_ai.Agent.instrument_all]
                 will be used, which defaults to False.
+                See the [Debugging and Monitoring guide](https://ai.pydantic.dev/logfire/) for more info.
         """
         if model is None or defer_model_check:
             self.model = model
@@ -215,6 +225,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)
@@ -435,7 +446,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits = usage_limits or _usage.UsageLimits()
 
         if isinstance(model_used, InstrumentedModel):
-            tracer = model_used.options.tracer
+            tracer = model_used.settings.tracer
         else:
             tracer = NoOpTracer()
         agent_name = self.name or 'agent'
@@ -461,6 +472,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 +939,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 +952,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 +994,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 +1003,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 +1018,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 +1035,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 +1048,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 +1090,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 +1124,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 +1136,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 +1287,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.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/common_tools/duckduckgo.py

@@ -13,7 +13,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `duckduckgo-search` to use the DuckDuckGo search tool, '
-        "you can use the `duckduckgo` optional group — `pip install 'pydantic-ai-slim[duckduckgo]'`"
+        'you can use the `duckduckgo` optional group — `pip install "pydantic-ai-slim[duckduckgo]"`'
     ) from _import_error
 
 __all__ = ('duckduckgo_search_tool',)

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/common_tools/tavily.py

@@ -11,7 +11,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `tavily-python` to use the Tavily search tool, '
-        "you can use the `tavily` optional group — `pip install 'pydantic-ai-slim[tavily]'`"
+        'you can use the `tavily` optional group — `pip install "pydantic-ai-slim[tavily]"`'
     ) from _import_error
 
 __all__ = ('tavily_search_tool',)

pydantic_ai_slim-0.0.43/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.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/messages.py

@@ -26,6 +26,9 @@ class SystemPromptPart:
     content: str
     """The content of the prompt."""
 
+    timestamp: datetime = field(default_factory=_now_utc)
+    """The timestamp of the prompt."""
+
     dynamic_ref: str | None = None
     """The ref of the dynamic system prompt function that generated this part.
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/anthropic.py

@@ -65,7 +65,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `anthropic` to use the Anthropic model, '
-        "you can use the `anthropic` optional group — `pip install 'pydantic-ai-slim[anthropic]'`"
+        'you can use the `anthropic` optional group — `pip install "pydantic-ai-slim[anthropic]"`'
     ) from _import_error
 
 LatestAnthropicModelNames = Literal[

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/cohere.py

@@ -50,7 +50,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `cohere` to use the Cohere model, '
-        "you can use the `cohere` optional group — `pip install 'pydantic-ai-slim[cohere]'`"
+        'you can use the `cohere` optional group — `pip install "pydantic-ai-slim[cohere]"`'
     ) from _import_error
 
 LatestCohereModelNames = Literal[

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/groq.py

@@ -41,7 +41,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `groq` to use the Groq model, '
-        "you can use the `groq` optional group — `pip install 'pydantic-ai-slim[groq]'`"
+        'you can use the `groq` optional group — `pip install "pydantic-ai-slim[groq]"`'
     ) from _import_error
 
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/instrumented.py

@@ -52,7 +52,9 @@ class InstrumentationSettings:
 
     - `Agent(instrument=...)`
     - [`Agent.instrument_all()`][pydantic_ai.agent.Agent.instrument_all]
-    - `InstrumentedModel`
+    - [`InstrumentedModel`][pydantic_ai.models.instrumented.InstrumentedModel]
+
+    See the [Debugging and Monitoring guide](https://ai.pydantic.dev/logfire/) for more info.
     """
 
     tracer: Tracer = field(repr=False)
@@ -94,9 +96,13 @@ GEN_AI_REQUEST_MODEL_ATTRIBUTE = 'gen_ai.request.model'
 
 @dataclass
 class InstrumentedModel(WrapperModel):
-    """Model which is instrumented with OpenTelemetry."""
+    """Model which wraps another model so that requests are instrumented with OpenTelemetry.
+
+    See the [Debugging and Monitoring guide](https://ai.pydantic.dev/logfire/) for more info.
+    """
 
-    options: InstrumentationSettings
+    settings: InstrumentationSettings
+    """Configuration for instrumenting requests."""
 
     def __init__(
         self,
@@ -104,7 +110,7 @@ class InstrumentedModel(WrapperModel):
         options: InstrumentationSettings | None = None,
     ) -> None:
         super().__init__(wrapped)
-        self.options = options or InstrumentationSettings()
+        self.settings = options or InstrumentationSettings()
 
     async def request(
         self,
@@ -156,7 +162,7 @@ class InstrumentedModel(WrapperModel):
                 if isinstance(value := model_settings.get(key), (float, int)):
                     attributes[f'gen_ai.request.{key}'] = value
 
-        with self.options.tracer.start_as_current_span(span_name, attributes=attributes) as span:
+        with self.settings.tracer.start_as_current_span(span_name, attributes=attributes) as span:
 
             def finish(response: ModelResponse, usage: Usage):
                 if not span.is_recording():
@@ -190,9 +196,9 @@ class InstrumentedModel(WrapperModel):
             yield finish
 
     def _emit_events(self, span: Span, events: list[Event]) -> None:
-        if self.options.event_mode == 'logs':
+        if self.settings.event_mode == 'logs':
             for event in events:
-                self.options.event_logger.emit(event)
+                self.settings.event_logger.emit(event)
         else:
             attr_name = 'events'
             span.set_attributes(

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/mistral.py

@@ -75,7 +75,7 @@ try:
 except ImportError as e:
     raise ImportError(
         'Please install `mistral` to use the Mistral model, '
-        "you can use the `mistral` optional group — `pip install 'pydantic-ai-slim[mistral]'`"
+        'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
     ) from e
 
 LatestMistralModelNames = Literal[

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/openai.py

@@ -57,7 +57,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `openai` to use the OpenAI model, '
-        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
     ) from _import_error
 
 OpenAIModelName = Union[str, ChatModel]

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/vertexai.py

@@ -27,7 +27,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install `google-auth` to use the VertexAI model, '
-        "you can use the `vertexai` optional group — `pip install 'pydantic-ai-slim[vertexai]'`"
+        'you can use the `vertexai` optional group — `pip install "pydantic-ai-slim[vertexai]"`'
     ) from _import_error
 
 VERTEX_AI_URL_TEMPLATE = (

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/models/wrapper.py

@@ -13,9 +13,13 @@ from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse, i
 
 @dataclass(init=False)
 class WrapperModel(Model):
-    """Model which wraps another model."""
+    """Model which wraps another model.
+
+    Does nothing on its own, used as a base class.
+    """
 
     wrapped: Model
+    """The underlying model being wrapped."""
 
     def __init__(self, wrapped: Model | KnownModelName):
         self.wrapped = infer_model(wrapped)

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/anthropic.py

@@ -12,7 +12,7 @@ try:
 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]'`"
+        'you can use the `anthropic` optional group — `pip install "pydantic-ai-slim[anthropic]"`'
     ) from _import_error
 
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/azure.py

@@ -13,7 +13,7 @@ try:
 except ImportError as _import_error:  # pragma: no cover
     raise ImportError(
         'Please install the `openai` package to use the Azure provider, '
-        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
     ) from _import_error
 
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/bedrock.py

@@ -11,7 +11,7 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install the `boto3` package to use the Bedrock provider, '
-        "you can use the `bedrock` optional group — `pip install 'pydantic-ai-slim[bedrock]'`"
+        'you can use the `bedrock` optional group — `pip install "pydantic-ai-slim[bedrock]"`'
     ) from _import_error
 
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/deepseek.py

@@ -13,7 +13,7 @@ try:
 except ImportError as _import_error:  # pragma: no cover
     raise ImportError(
         'Please install the `openai` package to use the DeepSeek provider, '
-        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
     ) from _import_error
 
 from . import Provider

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/google_vertex.py

@@ -2,7 +2,6 @@ from __future__ import annotations as _annotations
 
 import functools
 from collections.abc import AsyncGenerator, Mapping
-from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Literal, overload
 
@@ -22,15 +21,12 @@ try:
 except ImportError as _import_error:
     raise ImportError(
         'Please install the `google-auth` package to use the Google Vertex AI provider, '
-        "you can use the `vertexai` optional group — `pip install 'pydantic-ai-slim[vertexai]'`"
+        'you can use the `vertexai` optional group — `pip install "pydantic-ai-slim[vertexai]"`'
     ) from _import_error
 
 
 __all__ = ('GoogleVertexProvider',)
 
-# default expiry is 3600 seconds
-MAX_TOKEN_AGE = timedelta(seconds=3000)
-
 
 class GoogleVertexProvider(Provider[httpx.AsyncClient]):
     """Provider for Vertex AI API."""
@@ -131,19 +127,21 @@ class _VertexAIAuth(httpx.Auth):
         self.region = region
 
         self.credentials = None
-        self.token_created: datetime | None = None
 
     async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
         if self.credentials is None:
             self.credentials = await self._get_credentials()
-        if self.credentials.token is None or self._token_expired():  # type: ignore[reportUnknownMemberType]
-            await anyio.to_thread.run_sync(self._refresh_token)
-            self.token_created = datetime.now()
+        if self.credentials.token is None:  # type: ignore[reportUnknownMemberType]
+            await self._refresh_token()
         request.headers['Authorization'] = f'Bearer {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
-
         # NOTE: This workaround is in place because we might get the project_id from the credentials.
         request.url = httpx.URL(str(request.url).replace('projects/None', f'projects/{self.project_id}'))
-        yield request
+        response = yield request
+
+        if response.status_code == 401:
+            await self._refresh_token()
+            request.headers['Authorization'] = f'Bearer {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
+            yield request
 
     async def _get_credentials(self) -> BaseCredentials | ServiceAccountCredentials:
         if self.service_account_file is not None:
@@ -166,15 +164,9 @@ class _VertexAIAuth(httpx.Auth):
             self.project_id = creds_project_id
         return creds
 
-    def _token_expired(self) -> bool:
-        if self.token_created is None:
-            return True
-        else:
-            return (datetime.now() - self.token_created) > MAX_TOKEN_AGE
-
-    def _refresh_token(self) -> str:  # pragma: no cover
+    async def _refresh_token(self) -> str:  # pragma: no cover
         assert self.credentials is not None
-        self.credentials.refresh(Request())  # type: ignore[reportUnknownMemberType]
+        await anyio.to_thread.run_sync(self.credentials.refresh, Request())  # type: ignore[reportUnknownMemberType]
         assert isinstance(self.credentials.token, str), f'Expected token to be a string, got {self.credentials.token}'  # type: ignore[reportUnknownMemberType]
         return self.credentials.token
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/groq.py

@@ -12,7 +12,7 @@ try:
 except ImportError as _import_error:  # pragma: no cover
     raise ImportError(
         'Please install the `groq` package to use the Groq provider, '
-        "you can use the `groq` optional group — `pip install 'pydantic-ai-slim[groq]'`"
+        'you can use the `groq` optional group — `pip install "pydantic-ai-slim[groq]"`'
     ) from _import_error
 
 
@@ -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.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/mistral.py

@@ -12,7 +12,7 @@ try:
 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]'`"
+        'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
     ) from e
 
 
@@ -55,19 +55,19 @@ class MistralProvider(Provider[Mistral]):
             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.
         """
-        api_key = api_key or os.environ.get('MISTRAL_API_KEY')
-
-        if api_key is None and mistral_client is None:
-            raise ValueError(
-                'Set the `MISTRAL_API_KEY` environment variable or pass it via `MistralProvider(api_key=...)`'
-                'to use the Mistral provider.'
-            )
-
         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
-        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())
+            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.41 → pydantic_ai_slim-0.0.43}/pydantic_ai/providers/openai.py

@@ -11,7 +11,7 @@ try:
 except ImportError as _import_error:  # pragma: no cover
     raise ImportError(
         'Please install the `openai` package to use the OpenAI provider, '
-        "you can use the `openai` optional group — `pip install 'pydantic-ai-slim[openai]'`"
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
     ) from _import_error
 
 

{pydantic_ai_slim-0.0.41 → pydantic_ai_slim-0.0.43}/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.41 → pydantic_ai_slim-0.0.43}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.41"
+version = "0.0.43"
 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.41",
+    "pydantic-graph==0.0.43",
     "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"