pydantic-ai-slim 1.2.1__tar.gz → 1.4.0__tar.gz

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 1.2.1
+Version: 1.4.0
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Project-URL: Homepage, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
 Project-URL: Source, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
@@ -29,11 +29,11 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Requires-Dist: exceptiongroup; python_version < '3.11'
-Requires-Dist: genai-prices>=0.0.31
+Requires-Dist: genai-prices>=0.0.35
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==1.2.1
+Requires-Dist: pydantic-graph==1.4.0
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -57,9 +57,9 @@ Requires-Dist: dbos>=1.14.0; extra == 'dbos'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==1.2.1; extra == 'evals'
+Requires-Dist: pydantic-evals==1.4.0; extra == 'evals'
 Provides-Extra: google
-Requires-Dist: google-genai>=1.31.0; extra == 'google'
+Requires-Dist: google-genai>=1.46.0; extra == 'google'
 Provides-Extra: groq
 Requires-Dist: groq>=0.25.0; extra == 'groq'
 Provides-Extra: huggingface

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/__init__.py

@@ -12,6 +12,7 @@ from .agent import (
 from .builtin_tools import (
     CodeExecutionTool,
     ImageGenerationTool,
+    MCPServerTool,
     MemoryTool,
     UrlContextTool,
     WebSearchTool,
@@ -22,6 +23,7 @@ from .exceptions import (
     ApprovalRequired,
     CallDeferred,
     FallbackExceptionGroup,
+    IncompleteToolCall,
     ModelHTTPError,
     ModelRetry,
     UnexpectedModelBehavior,
@@ -124,6 +126,7 @@ __all__ = (
     'ModelRetry',
     'ModelHTTPError',
     'FallbackExceptionGroup',
+    'IncompleteToolCall',
     'UnexpectedModelBehavior',
     'UsageLimitExceeded',
     'UserError',
@@ -211,6 +214,7 @@ __all__ = (
     'CodeExecutionTool',
     'ImageGenerationTool',
     'MemoryTool',
+    'MCPServerTool',
     # output
     'ToolOutput',
     'NativeOutput',

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/_agent_graph.py

@@ -92,9 +92,28 @@ class GraphAgentState:
     retries: int = 0
     run_step: int = 0
 
-    def increment_retries(self, max_result_retries: int, error: BaseException | None = None) -> None:
+    def increment_retries(
+        self,
+        max_result_retries: int,
+        error: BaseException | None = None,
+        model_settings: ModelSettings | None = None,
+    ) -> None:
         self.retries += 1
         if self.retries > max_result_retries:
+            if (
+                self.message_history
+                and isinstance(model_response := self.message_history[-1], _messages.ModelResponse)
+                and model_response.finish_reason == 'length'
+                and model_response.parts
+                and isinstance(tool_call := model_response.parts[-1], _messages.ToolCallPart)
+            ):
+                try:
+                    tool_call.args_as_dict()
+                except Exception:
+                    max_tokens = (model_settings or {}).get('max_tokens') if model_settings else None
+                    raise exceptions.IncompleteToolCall(
+                        f'Model token limit ({max_tokens if max_tokens is not None else "provider default"}) exceeded while emitting a tool call, resulting in incomplete arguments. Increase max tokens or simplify tool call arguments to fit within limit.'
+                    )
             message = f'Exceeded maximum retries ({max_result_retries}) for output validation'
             if error:
                 if isinstance(error, exceptions.UnexpectedModelBehavior) and error.__cause__ is not None:
@@ -568,8 +587,12 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     # resubmit the most recent request that resulted in an empty response,
                     # as the empty response and request will not create any items in the API payload,
                     # in the hope the model will return a non-empty response this time.
-                    ctx.state.increment_retries(ctx.deps.max_result_retries)
-                    self._next_node = ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[]))
+                    ctx.state.increment_retries(ctx.deps.max_result_retries, model_settings=ctx.deps.model_settings)
+                    run_context = build_run_context(ctx)
+                    instructions = await ctx.deps.get_instructions(run_context)
+                    self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
+                        _messages.ModelRequest(parts=[], instructions=instructions)
+                    )
                     return
 
                 text = ''
@@ -630,8 +653,14 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     )
                     raise ToolRetryError(m)
                 except ToolRetryError as e:
-                    ctx.state.increment_retries(ctx.deps.max_result_retries, e)
-                    self._next_node = ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
+                    ctx.state.increment_retries(
+                        ctx.deps.max_result_retries, error=e, model_settings=ctx.deps.model_settings
+                    )
+                    run_context = build_run_context(ctx)
+                    instructions = await ctx.deps.get_instructions(run_context)
+                    self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
+                        _messages.ModelRequest(parts=[e.tool_retry], instructions=instructions)
+                    )
 
             self._events_iterator = _run_stream()
 
@@ -788,10 +817,14 @@ async def process_tool_calls(  # noqa: C901
             try:
                 result_data = await tool_manager.handle_call(call)
             except exceptions.UnexpectedModelBehavior as e:
-                ctx.state.increment_retries(ctx.deps.max_result_retries, e)
+                ctx.state.increment_retries(
+                    ctx.deps.max_result_retries, error=e, model_settings=ctx.deps.model_settings
+                )
                 raise e  # pragma: lax no cover
             except ToolRetryError as e:
-                ctx.state.increment_retries(ctx.deps.max_result_retries, e)
+                ctx.state.increment_retries(
+                    ctx.deps.max_result_retries, error=e, model_settings=ctx.deps.model_settings
+                )
                 yield _messages.FunctionToolCallEvent(call)
                 output_parts.append(e.tool_retry)
                 yield _messages.FunctionToolResultEvent(e.tool_retry)
@@ -820,7 +853,7 @@ async def process_tool_calls(  # noqa: C901
 
     # Then, we handle unknown tool calls
     if tool_calls_by_kind['unknown']:
-        ctx.state.increment_retries(ctx.deps.max_result_retries)
+        ctx.state.increment_retries(ctx.deps.max_result_retries, model_settings=ctx.deps.model_settings)
         calls_to_run.extend(tool_calls_by_kind['unknown'])
 
     calls_to_run_results: dict[str, DeferredToolResult] = {}

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/agent/__init__.py

@@ -542,6 +542,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
+
         model_used = self._get_model(model)
         del model
 
@@ -607,16 +608,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         else:
             instrumentation_settings = None
             tracer = NoOpTracer()
-        if builtin_tools:
-            # Deduplicate builtin tools passed to the agent and the run based on type
-            builtin_tools = list(
-                {
-                    **({type(tool): tool for tool in self._builtin_tools or []}),
-                    **({type(tool): tool for tool in builtin_tools}),
-                }.values()
-            )
-        else:
-            builtin_tools = list(self._builtin_tools)
+
         graph_deps = _agent_graph.GraphAgentDeps[AgentDepsT, RunOutputDataT](
             user_deps=deps,
             prompt=user_prompt,
@@ -629,7 +621,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             output_schema=output_schema,
             output_validators=output_validators,
             history_processors=self.history_processors,
-            builtin_tools=builtin_tools,
+            builtin_tools=[*self._builtin_tools, *(builtin_tools or [])],
             tool_manager=tool_manager,
             tracer=tracer,
             get_instructions=get_instructions,
@@ -662,14 +654,14 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         )
 
         try:
-            async with toolset:
-                async with graph.iter(
-                    start_node,
-                    state=state,
-                    deps=graph_deps,
-                    span=use_span(run_span) if run_span.is_recording() else None,
-                    infer_name=False,
-                ) as graph_run:
+            async with graph.iter(
+                start_node,
+                state=state,
+                deps=graph_deps,
+                span=use_span(run_span) if run_span.is_recording() else None,
+                infer_name=False,
+            ) as graph_run:
+                async with toolset:
                     agent_run = AgentRun(graph_run)
                     yield agent_run
                     if (final_result := agent_run.result) is not None and run_span.is_recording():

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/builtin_tools.py

@@ -2,13 +2,12 @@ from __future__ import annotations as _annotations
 
 from abc import ABC
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import Annotated, Any, Literal, Union
 
+import pydantic
+from pydantic_core import core_schema
 from typing_extensions import TypedDict
 
-if TYPE_CHECKING:
-    from .builtin_tools import AbstractBuiltinTool
-
 __all__ = (
     'AbstractBuiltinTool',
     'WebSearchTool',
@@ -17,8 +16,11 @@ __all__ = (
     'UrlContextTool',
     'ImageGenerationTool',
     'MemoryTool',
+    'MCPServerTool',
 )
 
+_BUILTIN_TOOL_TYPES: dict[str, type[AbstractBuiltinTool]] = {}
+
 
 @dataclass(kw_only=True)
 class AbstractBuiltinTool(ABC):
@@ -32,6 +34,34 @@ class AbstractBuiltinTool(ABC):
     kind: str = 'unknown_builtin_tool'
     """Built-in tool identifier, this should be available on all built-in tools as a discriminator."""
 
+    @property
+    def unique_id(self) -> str:
+        """A unique identifier for the builtin tool.
+
+        If multiple instances of the same builtin tool can be passed to the model, subclasses should override this property to allow them to be distinguished.
+        """
+        return self.kind
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        _BUILTIN_TOOL_TYPES[cls.kind] = cls
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, _source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        if cls is not AbstractBuiltinTool:
+            return handler(cls)
+
+        tools = _BUILTIN_TOOL_TYPES.values()
+        if len(tools) == 1:  # pragma: no cover
+            tools_type = next(iter(tools))
+        else:
+            tools_annotated = [Annotated[tool, pydantic.Tag(tool.kind)] for tool in tools]
+            tools_type = Annotated[Union[tuple(tools_annotated)], pydantic.Discriminator(_tool_discriminator)]  # noqa: UP007
+
+        return handler(tools_type)
+
 
 @dataclass(kw_only=True)
 class WebSearchTool(AbstractBuiltinTool):
@@ -120,6 +150,7 @@ class WebSearchUserLocation(TypedDict, total=False):
     """The timezone of the user's location."""
 
 
+@dataclass(kw_only=True)
 class CodeExecutionTool(AbstractBuiltinTool):
     """A builtin tool that allows your agent to execute code.
 
@@ -134,6 +165,7 @@ class CodeExecutionTool(AbstractBuiltinTool):
     """The kind of tool."""
 
 
+@dataclass(kw_only=True)
 class UrlContextTool(AbstractBuiltinTool):
     """Allows your agent to access contents from URLs.
 
@@ -227,6 +259,7 @@ class ImageGenerationTool(AbstractBuiltinTool):
     """The kind of tool."""
 
 
+@dataclass(kw_only=True)
 class MemoryTool(AbstractBuiltinTool):
     """A builtin tool that allows your agent to use memory.
 
@@ -237,3 +270,72 @@ class MemoryTool(AbstractBuiltinTool):
 
     kind: str = 'memory'
     """The kind of tool."""
+
+
+@dataclass(kw_only=True)
+class MCPServerTool(AbstractBuiltinTool):
+    """A builtin tool that allows your agent to use MCP servers.
+
+    Supported by:
+
+    * OpenAI Responses
+    * Anthropic
+    """
+
+    id: str
+    """A unique identifier for the MCP server."""
+
+    url: str
+    """The URL of the MCP server to use.
+
+    For OpenAI Responses, it is possible to use `connector_id` by providing it as `x-openai-connector:<connector_id>`.
+    """
+
+    authorization_token: str | None = None
+    """Authorization header to use when making requests to the MCP server.
+
+    Supported by:
+
+    * OpenAI Responses
+    * Anthropic
+    """
+
+    description: str | None = None
+    """A description of the MCP server.
+
+    Supported by:
+
+    * OpenAI Responses
+    """
+
+    allowed_tools: list[str] | None = None
+    """A list of tools that the MCP server can use.
+
+    Supported by:
+
+    * OpenAI Responses
+    * Anthropic
+    """
+
+    headers: dict[str, str] | None = None
+    """Optional HTTP headers to send to the MCP server.
+
+    Use for authentication or other purposes.
+
+    Supported by:
+
+    * OpenAI Responses
+    """
+
+    kind: str = 'mcp_server'
+
+    @property
+    def unique_id(self) -> str:
+        return ':'.join([self.kind, self.id])
+
+
+def _tool_discriminator(tool_data: dict[str, Any] | AbstractBuiltinTool) -> str:
+    if isinstance(tool_data, dict):
+        return tool_data.get('kind', AbstractBuiltinTool.kind)
+    else:
+        return tool_data.kind

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/exceptions.py

@@ -23,6 +23,7 @@ __all__ = (
     'UnexpectedModelBehavior',
     'UsageLimitExceeded',
     'ModelHTTPError',
+    'IncompleteToolCall',
     'FallbackExceptionGroup',
 )
 
@@ -168,3 +169,7 @@ class ToolRetryError(Exception):
     def __init__(self, tool_retry: RetryPromptPart):
         self.tool_retry = tool_retry
         super().__init__()
+
+
+class IncompleteToolCall(UnexpectedModelBehavior):
+    """Error raised when a model stops due to token limit while emitting a tool call."""

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/mcp.py

@@ -441,14 +441,9 @@ class MCPServerStdio(MCPServer):
         'uv', args=['run', 'mcp-run-python', 'stdio'], timeout=10
     )
     agent = Agent('openai:gpt-4o', toolsets=[server])
-
-    async def main():
-        async with agent:  # (2)!
-            ...
     ```
 
     1. See [MCP Run Python](https://github.com/pydantic/mcp-run-python) for more information.
-    2. This will start the server as a subprocess and connect to it.
     """
 
     command: str
@@ -788,13 +783,7 @@ class MCPServerSSE(_MCPServerHTTP):
 
     server = MCPServerSSE('http://localhost:3001/sse')
     agent = Agent('openai:gpt-4o', toolsets=[server])
-
-    async def main():
-        async with agent:  # (1)!
-            ...
     ```
-
-    1. This will connect to a server running on `localhost:3001`.
     """
 
     @classmethod
@@ -837,13 +826,7 @@ class MCPServerHTTP(MCPServerSSE):
 
     server = MCPServerHTTP('http://localhost:3001/sse')
     agent = Agent('openai:gpt-4o', toolsets=[server])
-
-    async def main():
-        async with agent:  # (2)!
-            ...
     ```
-
-    1. This will connect to a server running on `localhost:3001`.
     """
 
 
@@ -862,12 +845,8 @@ class MCPServerStreamableHTTP(_MCPServerHTTP):
     from pydantic_ai import Agent
     from pydantic_ai.mcp import MCPServerStreamableHTTP
 
-    server = MCPServerStreamableHTTP('http://localhost:8000/mcp')  # (1)!
+    server = MCPServerStreamableHTTP('http://localhost:8000/mcp')
     agent = Agent('openai:gpt-4o', toolsets=[server])
-
-    async def main():
-        async with agent:  # (2)!
-            ...
     ```
     """
 

{pydantic_ai_slim-1.2.1 → pydantic_ai_slim-1.4.0}/pydantic_ai/models/__init__.py

@@ -43,6 +43,7 @@ from ..messages import (
 )
 from ..output import OutputMode
 from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
+from ..providers import infer_provider
 from ..settings import ModelSettings, merge_model_settings
 from ..tools import ToolDefinition
 from ..usage import RequestUsage
@@ -129,15 +130,8 @@ KnownModelName = TypeAliasType(
         'cerebras:qwen-3-235b-a22b-thinking-2507',
         'cohere:c4ai-aya-expanse-32b',
         'cohere:c4ai-aya-expanse-8b',
-        'cohere:command',
-        'cohere:command-light',
-        'cohere:command-light-nightly',
         'cohere:command-nightly',
-        'cohere:command-r',
-        'cohere:command-r-03-2024',
         'cohere:command-r-08-2024',
-        'cohere:command-r-plus',
-        'cohere:command-r-plus-04-2024',
         'cohere:command-r-plus-08-2024',
         'cohere:command-r7b-12-2024',
         'deepseek:deepseek-chat',
@@ -416,9 +410,17 @@ class Model(ABC):
         they need to customize the preparation flow further, but most implementations should simply call
         ``self.prepare_request(...)`` at the start of their ``request`` (and related) methods.
         """
-        merged_settings = merge_model_settings(self.settings, model_settings)
-        customized_parameters = self.customize_request_parameters(model_request_parameters)
-        return merged_settings, customized_parameters
+        model_settings = merge_model_settings(self.settings, model_settings)
+
+        if builtin_tools := model_request_parameters.builtin_tools:
+            # Deduplicate builtin tools
+            model_request_parameters = replace(
+                model_request_parameters,
+                builtin_tools=list({tool.unique_id: tool for tool in builtin_tools}.values()),
+            )
+
+        model_request_parameters = self.customize_request_parameters(model_request_parameters)
+        return model_settings, model_request_parameters
 
     @property
     @abstractmethod
@@ -644,41 +646,39 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         return TestModel()
 
     try:
-        provider, model_name = model.split(':', maxsplit=1)
+        provider_name, model_name = model.split(':', maxsplit=1)
     except ValueError:
-        provider = None
+        provider_name = None
         model_name = model
         if model_name.startswith(('gpt', 'o1', 'o3')):
-            provider = 'openai'
+            provider_name = 'openai'
         elif model_name.startswith('claude'):
-            provider = 'anthropic'
+            provider_name = 'anthropic'
         elif model_name.startswith('gemini'):
-            provider = 'google-gla'
+            provider_name = 'google-gla'
 
-        if provider is not None:
+        if provider_name is not None:
             warnings.warn(
-                f"Specifying a model name without a provider prefix is deprecated. Instead of {model_name!r}, use '{provider}:{model_name}'.",
+                f"Specifying a model name without a provider prefix is deprecated. Instead of {model_name!r}, use '{provider_name}:{model_name}'.",
                 DeprecationWarning,
             )
         else:
             raise UserError(f'Unknown model: {model}')
 
-    if provider == 'vertexai':  # pragma: no cover
+    if provider_name == 'vertexai':  # pragma: no cover
         warnings.warn(
             "The 'vertexai' provider name is deprecated. Use 'google-vertex' instead.",
             DeprecationWarning,
         )
-        provider = 'google-vertex'
-
-    if provider == 'gateway':
-        from ..providers.gateway import infer_model as infer_model_from_gateway
+        provider_name = 'google-vertex'
 
-        return infer_model_from_gateway(model_name)
-    elif provider == 'cohere':
-        from .cohere import CohereModel
+    provider = infer_provider(provider_name)
 
-        return CohereModel(model_name, provider=provider)
-    elif provider in (
+    model_kind = provider_name
+    if model_kind.startswith('gateway/'):
+        model_kind = provider_name.removeprefix('gateway/')
+    if model_kind in (
+        'openai',
         'azure',
         'deepseek',
         'cerebras',
@@ -688,42 +688,50 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'heroku',
         'moonshotai',
         'ollama',
-        'openai',
-        'openai-chat',
         'openrouter',
         'together',
         'vercel',
         'litellm',
         'nebius',
+        'ovhcloud',
     ):
+        model_kind = 'openai-chat'
+    elif model_kind in ('google-gla', 'google-vertex'):
+        model_kind = 'google'
+
+    if model_kind == 'openai-chat':
         from .openai import OpenAIChatModel
 
         return OpenAIChatModel(model_name, provider=provider)
-    elif provider == 'openai-responses':
+    elif model_kind == 'openai-responses':
         from .openai import OpenAIResponsesModel
 
-        return OpenAIResponsesModel(model_name, provider='openai')
-    elif provider in ('google-gla', 'google-vertex'):
+        return OpenAIResponsesModel(model_name, provider=provider)
+    elif model_kind == 'google':
         from .google import GoogleModel
 
         return GoogleModel(model_name, provider=provider)
-    elif provider == 'groq':
+    elif model_kind == 'groq':
         from .groq import GroqModel
 
         return GroqModel(model_name, provider=provider)
-    elif provider == 'mistral':
+    elif model_kind == 'cohere':
+        from .cohere import CohereModel
+
+        return CohereModel(model_name, provider=provider)
+    elif model_kind == 'mistral':
         from .mistral import MistralModel
 
         return MistralModel(model_name, provider=provider)
-    elif provider == 'anthropic':
+    elif model_kind == 'anthropic':
         from .anthropic import AnthropicModel
 
         return AnthropicModel(model_name, provider=provider)
-    elif provider == 'bedrock':
+    elif model_kind == 'bedrock':
         from .bedrock import BedrockConverseModel
 
         return BedrockConverseModel(model_name, provider=provider)
-    elif provider == 'huggingface':
+    elif model_kind == 'huggingface':
         from .huggingface import HuggingFaceModel
 
         return HuggingFaceModel(model_name, provider=provider)