pydantic-ai-slim 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

pydantic_ai/common_tools/duckduckgo.py

@@ -9,10 +9,13 @@ from typing_extensions import TypedDict
 from pydantic_ai.tools import Tool
 
 try:
-    from duckduckgo_search import DDGS
+    try:
+        from ddgs import DDGS
+    except ImportError:  # Fallback for older versions of ddgs
+        from duckduckgo_search import DDGS
 except ImportError as _import_error:
     raise ImportError(
-        'Please install `duckduckgo-search` to use the DuckDuckGo search tool, '
+        'Please install `ddgs` to use the DuckDuckGo search tool, '
         'you can use the `duckduckgo` optional group — `pip install "pydantic-ai-slim[duckduckgo]"`'
     ) from _import_error
 

pydantic_ai/exceptions.py

@@ -2,11 +2,15 @@ from __future__ import annotations as _annotations
 
 import json
 import sys
+from typing import TYPE_CHECKING
 
 if sys.version_info < (3, 11):
-    from exceptiongroup import ExceptionGroup  # pragma: lax no cover
+    from exceptiongroup import ExceptionGroup
 else:
-    ExceptionGroup = ExceptionGroup  # pragma: lax no cover
+    ExceptionGroup = ExceptionGroup
+
+if TYPE_CHECKING:
+    from .messages import RetryPromptPart
 
 __all__ = (
     'ModelRetry',
@@ -113,3 +117,11 @@ class ModelHTTPError(AgentRunError):
 
 class FallbackExceptionGroup(ExceptionGroup):
     """A group of exceptions that can be raised when all fallback models fail."""
+
+
+class ToolRetryError(Exception):
+    """Exception used to signal a `ToolRetry` message should be returned to the LLM."""
+
+    def __init__(self, tool_retry: RetryPromptPart):
+        self.tool_retry = tool_retry
+        super().__init__()

pydantic_ai/ext/aci.py

@@ -4,11 +4,13 @@ try:
 except ImportError as _import_error:
     raise ImportError('Please install `aci-sdk` to use ACI.dev tools') from _import_error
 
+from collections.abc import Sequence
 from typing import Any
 
 from aci import ACI
 
-from pydantic_ai import Tool
+from pydantic_ai.tools import Tool
+from pydantic_ai.toolsets.function import FunctionToolset
 
 
 def _clean_schema(schema):
@@ -22,10 +24,10 @@ def _clean_schema(schema):
 
 
 def tool_from_aci(aci_function: str, linked_account_owner_id: str) -> Tool:
-    """Creates a Pydantic AI tool proxy from an ACI function.
+    """Creates a Pydantic AI tool proxy from an ACI.dev function.
 
     Args:
-        aci_function: The ACI function to wrao.
+        aci_function: The ACI.dev function to wrap.
         linked_account_owner_id: The ACI user ID to execute the function on behalf of.
 
     Returns:
@@ -64,3 +66,10 @@ def tool_from_aci(aci_function: str, linked_account_owner_id: str) -> Tool:
         description=function_description,
         json_schema=json_schema,
     )
+
+
+class ACIToolset(FunctionToolset):
+    """A toolset that wraps ACI.dev tools."""
+
+    def __init__(self, aci_functions: Sequence[str], linked_account_owner_id: str):
+        super().__init__([tool_from_aci(aci_function, linked_account_owner_id) for aci_function in aci_functions])

pydantic_ai/ext/langchain.py

@@ -3,6 +3,7 @@ from typing import Any, Protocol
 from pydantic.json_schema import JsonSchemaValue
 
 from pydantic_ai.tools import Tool
+from pydantic_ai.toolsets.function import FunctionToolset
 
 
 class LangChainTool(Protocol):
@@ -23,7 +24,7 @@ class LangChainTool(Protocol):
     def run(self, *args: Any, **kwargs: Any) -> str: ...
 
 
-__all__ = ('tool_from_langchain',)
+__all__ = ('tool_from_langchain', 'LangChainToolset')
 
 
 def tool_from_langchain(langchain_tool: LangChainTool) -> Tool:
@@ -59,3 +60,10 @@ def tool_from_langchain(langchain_tool: LangChainTool) -> Tool:
         description=function_description,
         json_schema=schema,
     )
+
+
+class LangChainToolset(FunctionToolset):
+    """A toolset that wraps LangChain tools."""
+
+    def __init__(self, tools: list[LangChainTool]):
+        super().__init__([tool_from_langchain(tool) for tool in tools])

pydantic_ai/mcp.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 import base64
 import functools
 from abc import ABC, abstractmethod
+from asyncio import Lock
 from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field, replace
 from pathlib import Path
-from types import TracebackType
 from typing import Any, Callable
 
 import anyio
@@ -16,6 +16,11 @@ import pydantic_core
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from typing_extensions import Self, assert_never, deprecated
 
+from pydantic_ai._run_context import RunContext
+from pydantic_ai.tools import ToolDefinition
+
+from .toolsets.abstract import AbstractToolset, ToolsetTool
+
 try:
     from mcp import types as mcp_types
     from mcp.client.session import ClientSession, LoggingFnT
@@ -32,12 +37,18 @@ except ImportError as _import_error:
     ) from _import_error
 
 # after mcp imports so any import error maps to this file, not _mcp.py
-from . import _mcp, exceptions, messages, models, tools
+from . import _mcp, exceptions, messages, models
 
 __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP', 'MCPServerSSE', 'MCPServerStreamableHTTP'
 
+TOOL_SCHEMA_VALIDATOR = pydantic_core.SchemaValidator(
+    schema=pydantic_core.core_schema.dict_schema(
+        pydantic_core.core_schema.str_schema(), pydantic_core.core_schema.any_schema()
+    )
+)
+
 
-class MCPServer(ABC):
+class MCPServer(AbstractToolset[Any], ABC):
     """Base class for attaching agents to MCP servers.
 
     See <https://modelcontextprotocol.io> for more information.
@@ -50,15 +61,22 @@ class MCPServer(ABC):
     timeout: float = 5
     process_tool_call: ProcessToolCallback | None = None
     allow_sampling: bool = True
+    max_retries: int = 1
+    sampling_model: models.Model | None = None
     # } end of "abstract fields"
 
-    _running_count: int = 0
+    _enter_lock: Lock = field(compare=False)
+    _running_count: int
+    _exit_stack: AsyncExitStack | None
 
     _client: ClientSession
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
-    _exit_stack: AsyncExitStack
-    sampling_model: models.Model | None = None
+
+    def __post_init__(self):
+        self._enter_lock = Lock()
+        self._running_count = 0
+        self._exit_stack = None
 
     @abstractmethod
     @asynccontextmanager
@@ -74,47 +92,36 @@ class MCPServer(ABC):
         raise NotImplementedError('MCP Server subclasses must implement this method.')
         yield
 
-    def get_prefixed_tool_name(self, tool_name: str) -> str:
-        """Get the tool name with prefix if `tool_prefix` is set."""
-        return f'{self.tool_prefix}_{tool_name}' if self.tool_prefix else tool_name
-
-    def get_unprefixed_tool_name(self, tool_name: str) -> str:
-        """Get original tool name without prefix for calling tools."""
-        return tool_name.removeprefix(f'{self.tool_prefix}_') if self.tool_prefix else tool_name
+    @property
+    def name(self) -> str:
+        return repr(self)
 
     @property
-    def is_running(self) -> bool:
-        """Check if the MCP server is running."""
-        return bool(self._running_count)
+    def tool_name_conflict_hint(self) -> str:
+        return 'Consider setting `tool_prefix` to avoid name conflicts.'
 
-    async def list_tools(self) -> list[tools.ToolDefinition]:
+    async def list_tools(self) -> list[mcp_types.Tool]:
         """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.
         """
-        mcp_tools = await self._client.list_tools()
-        return [
-            tools.ToolDefinition(
-                name=self.get_prefixed_tool_name(tool.name),
-                description=tool.description,
-                parameters_json_schema=tool.inputSchema,
-            )
-            for tool in mcp_tools.tools
-        ]
+        async with self:  # Ensure server is running
+            result = await self._client.list_tools()
+        return result.tools
 
-    async def call_tool(
+    async def direct_call_tool(
         self,
-        tool_name: str,
-        arguments: dict[str, Any],
+        name: str,
+        args: dict[str, Any],
         metadata: dict[str, Any] | None = None,
     ) -> ToolResult:
         """Call a tool on the server.
 
         Args:
-            tool_name: The name of the tool to call.
-            arguments: The arguments to pass to the tool.
+            name: The name of the tool to call.
+            args: The arguments to pass to the tool.
             metadata: Request-level metadata (optional)
 
         Returns:
@@ -123,23 +130,23 @@ class MCPServer(ABC):
         Raises:
             ModelRetry: If the tool call fails.
         """
-        try:
-            # meta param is not provided by session yet, so build and can send_request directly.
-            result = await self._client.send_request(
-                mcp_types.ClientRequest(
-                    mcp_types.CallToolRequest(
-                        method='tools/call',
-                        params=mcp_types.CallToolRequestParams(
-                            name=self.get_unprefixed_tool_name(tool_name),
-                            arguments=arguments,
-                            _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
-                        ),
-                    )
-                ),
-                mcp_types.CallToolResult,
-            )
-        except McpError as e:
-            raise exceptions.ModelRetry(e.error.message)
+        async with self:  # Ensure server is running
+            try:
+                result = await self._client.send_request(
+                    mcp_types.ClientRequest(
+                        mcp_types.CallToolRequest(
+                            method='tools/call',
+                            params=mcp_types.CallToolRequestParams(
+                                name=name,
+                                arguments=args,
+                                _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
+                            ),
+                        )
+                    ),
+                    mcp_types.CallToolResult,
+                )
+            except McpError as e:
+                raise exceptions.ModelRetry(e.error.message)
 
         content = [self._map_tool_result_part(part) for part in result.content]
 
@@ -149,36 +156,80 @@ class MCPServer(ABC):
         else:
             return content[0] if len(content) == 1 else content
 
-    async def __aenter__(self) -> Self:
-        if self._running_count == 0:
-            self._exit_stack = AsyncExitStack()
-
-            self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(self.client_streams())
-            client = ClientSession(
-                read_stream=self._read_stream,
-                write_stream=self._write_stream,
-                sampling_callback=self._sampling_callback if self.allow_sampling else None,
-                logging_callback=self.log_handler,
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[Any],
+        tool: ToolsetTool[Any],
+    ) -> ToolResult:
+        if self.tool_prefix:
+            name = name.removeprefix(f'{self.tool_prefix}_')
+            ctx = replace(ctx, tool_name=name)
+
+        if self.process_tool_call is not None:
+            return await self.process_tool_call(ctx, self.direct_call_tool, name, tool_args)
+        else:
+            return await self.direct_call_tool(name, tool_args)
+
+    async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
+        return {
+            name: ToolsetTool(
+                toolset=self,
+                tool_def=ToolDefinition(
+                    name=name,
+                    description=mcp_tool.description,
+                    parameters_json_schema=mcp_tool.inputSchema,
+                ),
+                max_retries=self.max_retries,
+                args_validator=TOOL_SCHEMA_VALIDATOR,
             )
-            self._client = await self._exit_stack.enter_async_context(client)
+            for mcp_tool in await self.list_tools()
+            if (name := f'{self.tool_prefix}_{mcp_tool.name}' if self.tool_prefix else mcp_tool.name)
+        }
+
+    async def __aenter__(self) -> Self:
+        """Enter the MCP server context.
+
+        This will initialize the connection to the server.
+        If this server is an [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio], the server will first be started as a subprocess.
 
-            with anyio.fail_after(self.timeout):
-                await self._client.initialize()
+        This is a no-op if the MCP server has already been entered.
+        """
+        async with self._enter_lock:
+            if self._running_count == 0:
+                self._exit_stack = AsyncExitStack()
+
+                self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(
+                    self.client_streams()
+                )
+                client = ClientSession(
+                    read_stream=self._read_stream,
+                    write_stream=self._write_stream,
+                    sampling_callback=self._sampling_callback if self.allow_sampling else None,
+                    logging_callback=self.log_handler,
+                )
+                self._client = await self._exit_stack.enter_async_context(client)
+
+                with anyio.fail_after(self.timeout):
+                    await self._client.initialize()
 
-                if log_level := self.log_level:
-                    await self._client.set_logging_level(log_level)
-        self._running_count += 1
+                    if log_level := self.log_level:
+                        await self._client.set_logging_level(log_level)
+            self._running_count += 1
         return self
 
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> bool | None:
-        self._running_count -= 1
-        if self._running_count <= 0:
-            await self._exit_stack.aclose()
+    async def __aexit__(self, *args: Any) -> bool | None:
+        async with self._enter_lock:
+            self._running_count -= 1
+            if self._running_count == 0 and self._exit_stack is not None:
+                await self._exit_stack.aclose()
+                self._exit_stack = None
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the MCP server is running."""
+        return bool(self._running_count)
 
     async def _sampling_callback(
         self, context: RequestContext[ClientSession, Any], params: mcp_types.CreateMessageRequestParams
@@ -271,10 +322,10 @@ class MCPServerStdio(MCPServer):
             'stdio',
         ]
     )
-    agent = Agent('openai:gpt-4o', mcp_servers=[server])
+    agent = Agent('openai:gpt-4o', toolsets=[server])
 
     async def main():
-        async with agent.run_mcp_servers():  # (2)!
+        async with agent:  # (2)!
             ...
     ```
 
@@ -327,6 +378,12 @@ class MCPServerStdio(MCPServer):
     allow_sampling: bool = True
     """Whether to allow MCP sampling through this client."""
 
+    max_retries: int = 1
+    """The maximum number of times to retry a tool call."""
+
+    sampling_model: models.Model | None = None
+    """The model to use for sampling."""
+
     @asynccontextmanager
     async def client_streams(
         self,
@@ -422,6 +479,12 @@ class _MCPServerHTTP(MCPServer):
     allow_sampling: bool = True
     """Whether to allow MCP sampling through this client."""
 
+    max_retries: int = 1
+    """The maximum number of times to retry a tool call."""
+
+    sampling_model: models.Model | None = None
+    """The model to use for sampling."""
+
     @property
     @abstractmethod
     def _transport_client(
@@ -503,10 +566,10 @@ class MCPServerSSE(_MCPServerHTTP):
     from pydantic_ai.mcp import MCPServerSSE
 
     server = MCPServerSSE('http://localhost:3001/sse')  # (1)!
-    agent = Agent('openai:gpt-4o', mcp_servers=[server])
+    agent = Agent('openai:gpt-4o', toolsets=[server])
 
     async def main():
-        async with agent.run_mcp_servers():  # (2)!
+        async with agent:  # (2)!
             ...
     ```
 
@@ -537,10 +600,10 @@ class MCPServerHTTP(MCPServerSSE):
     from pydantic_ai.mcp import MCPServerHTTP
 
     server = MCPServerHTTP('http://localhost:3001/sse')  # (1)!
-    agent = Agent('openai:gpt-4o', mcp_servers=[server])
+    agent = Agent('openai:gpt-4o', toolsets=[server])
 
     async def main():
-        async with agent.run_mcp_servers():  # (2)!
+        async with agent:  # (2)!
             ...
     ```
 
@@ -566,10 +629,10 @@ class MCPServerStreamableHTTP(_MCPServerHTTP):
     from pydantic_ai.mcp import MCPServerStreamableHTTP
 
     server = MCPServerStreamableHTTP('http://localhost:8000/mcp')  # (1)!
-    agent = Agent('openai:gpt-4o', mcp_servers=[server])
+    agent = Agent('openai:gpt-4o', toolsets=[server])
 
     async def main():
-        async with agent.run_mcp_servers():  # (2)!
+        async with agent:  # (2)!
             ...
     ```
     """
@@ -586,14 +649,14 @@ ToolResult = (
     | list[Any]
     | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]]
 )
-"""The result type of a tool call."""
+"""The result type of an MCP tool call."""
 
 CallToolFunc = Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[ToolResult]]
 """A function type that represents a tool call."""
 
 ProcessToolCallback = Callable[
     [
-        tools.RunContext[Any],
+        RunContext[Any],
         CallToolFunc,
         str,
         dict[str, Any],

pydantic_ai/messages.py

@@ -282,6 +282,14 @@ class BinaryContent:
     media_type: AudioMediaType | ImageMediaType | DocumentMediaType | str
     """The media type of the binary data."""
 
+    identifier: str | None = None
+    """Identifier for the binary content, such as a URL or unique ID.
+
+    This identifier can be provided to the model in a message to allow it to refer to this file in a tool call argument, and the tool can look up the file in question by iterating over the message history and finding the matching `BinaryContent`.
+
+    This identifier is only automatically passed to the model when the `BinaryContent` is returned by a tool. If you're passing the `BinaryContent` as a user message, it's up to you to include a separate text part with the identifier, e.g. "This is file <identifier>:" preceding the `BinaryContent`.
+    """
+
     vendor_metadata: dict[str, Any] | None = None
     """Vendor-specific metadata for the file.
 
@@ -411,9 +419,9 @@ class UserPromptPart:
     """Part type identifier, this is available on all parts as a discriminator."""
 
     def otel_event(self, settings: InstrumentationSettings) -> Event:
-        content: str | list[dict[str, Any] | str]
+        content: str | list[dict[str, Any] | str] | dict[str, Any]
         if isinstance(self.content, str):
-            content = self.content
+            content = self.content if settings.include_content else {'kind': 'text'}
         else:
             content = []
             for part in self.content:
@@ -433,7 +441,9 @@ class UserPromptPart:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-tool_return_ta: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter(Any, config=pydantic.ConfigDict(defer_build=True))
+tool_return_ta: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter(
+    Any, config=pydantic.ConfigDict(defer_build=True, ser_json_bytes='base64', val_json_bytes='base64')
+)
 
 
 @dataclass(repr=False)
@@ -519,7 +529,7 @@ class RetryPromptPart:
     tool_call_id: str = field(default_factory=_generate_tool_call_id)
     """The tool call identifier, this is used by some models including OpenAI.
 
-    In case the tool call id is not provided by the model, PydanticAI will generate a random one.
+    In case the tool call id is not provided by the model, Pydantic AI will generate a random one.
     """
 
     timestamp: datetime = field(default_factory=_now_utc)
@@ -560,12 +570,12 @@ class RetryPromptPart:
 ModelRequestPart = Annotated[
     Union[SystemPromptPart, UserPromptPart, ToolReturnPart, RetryPromptPart], pydantic.Discriminator('part_kind')
 ]
-"""A message part sent by PydanticAI to a model."""
+"""A message part sent by Pydantic AI to a model."""
 
 
 @dataclass(repr=False)
 class ModelRequest:
-    """A request generated by PydanticAI and sent to a model, e.g. a message from the PydanticAI app to the model."""
+    """A request generated by Pydantic AI and sent to a model, e.g. a message from the Pydantic AI app to the model."""
 
     parts: list[ModelRequestPart]
     """The parts of the user message."""
@@ -643,7 +653,7 @@ class ToolCallPart:
     tool_call_id: str = field(default_factory=_generate_tool_call_id)
     """The tool call identifier, this is used by some models including OpenAI.
 
-    In case the tool call id is not provided by the model, PydanticAI will generate a random one.
+    In case the tool call id is not provided by the model, Pydantic AI will generate a random one.
     """
 
     part_kind: Literal['tool-call'] = 'tool-call'
@@ -691,7 +701,7 @@ ModelResponsePart = Annotated[Union[TextPart, ToolCallPart, ThinkingPart], pydan
 
 @dataclass(repr=False)
 class ModelResponse:
-    """A response from a model, e.g. a message from the model to the PydanticAI app."""
+    """A response from a model, e.g. a message from the model to the Pydantic AI app."""
 
     parts: list[ModelResponsePart]
     """The parts of the model message."""
@@ -743,7 +753,7 @@ class ModelResponse:
                         'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
                         'function': {
                             'name': part.tool_name,
-                            'arguments': part.args,
+                            **({'arguments': part.args} if settings.include_content else {}),
                         },
                     }
                 )