pydantic-ai-slim 0.2.19__tar.gz → 0.2.20__tar.gz

{pydantic_ai_slim-0.2.19 → pydantic_ai_slim-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.2.19
+Version: 0.2.20
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
 License-Expression: MIT
@@ -30,11 +30,11 @@ 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.2.19
+Requires-Dist: pydantic-graph==0.2.20
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
-Requires-Dist: fasta2a==0.2.19; extra == 'a2a'
+Requires-Dist: fasta2a==0.2.20; extra == 'a2a'
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.52.0; extra == 'anthropic'
 Provides-Extra: bedrock
@@ -48,7 +48,7 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.2.19; extra == 'evals'
+Requires-Dist: pydantic-evals==0.2.20; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.15.0; extra == 'google'
 Provides-Extra: groq

{pydantic_ai_slim-0.2.19 → pydantic_ai_slim-0.2.20}/pydantic_ai/_agent_graph.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
 from opentelemetry.trace import Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
 
+from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx  # type: ignore
 from pydantic_ai._utils import is_async_callable, run_in_executor
 from pydantic_graph import BaseNode, Graph, GraphRunContext
 from pydantic_graph.nodes import End, NodeRunEndT
@@ -50,8 +51,20 @@ OutputT = TypeVar('OutputT')
 
 _HistoryProcessorSync = Callable[[list[_messages.ModelMessage]], list[_messages.ModelMessage]]
 _HistoryProcessorAsync = Callable[[list[_messages.ModelMessage]], Awaitable[list[_messages.ModelMessage]]]
-HistoryProcessor = Union[_HistoryProcessorSync, _HistoryProcessorAsync]
-"""A function that processes a list of model messages and returns a list of model messages."""
+_HistoryProcessorSyncWithCtx = Callable[[RunContext[DepsT], list[_messages.ModelMessage]], list[_messages.ModelMessage]]
+_HistoryProcessorAsyncWithCtx = Callable[
+    [RunContext[DepsT], list[_messages.ModelMessage]], Awaitable[list[_messages.ModelMessage]]
+]
+HistoryProcessor = Union[
+    _HistoryProcessorSync,
+    _HistoryProcessorAsync,
+    _HistoryProcessorSyncWithCtx[DepsT],
+    _HistoryProcessorAsyncWithCtx[DepsT],
+]
+"""A function that processes a list of model messages and returns a list of model messages.
+
+Can optionally accept a `RunContext` as a parameter.
+"""
 
 
 @dataclasses.dataclass
@@ -92,7 +105,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     output_schema: _output.OutputSchema[OutputDataT] | None
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
-    history_processors: Sequence[HistoryProcessor]
+    history_processors: Sequence[HistoryProcessor[DepsT]]
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
     mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
@@ -328,7 +341,9 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
-        message_history = await _process_message_history(ctx.state.message_history, ctx.deps.history_processors)
+        message_history = await _process_message_history(
+            ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
+        )
         async with ctx.deps.model.request_stream(
             message_history, model_settings, model_request_parameters
         ) as streamed_response:
@@ -352,7 +367,9 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
-        message_history = await _process_message_history(ctx.state.message_history, ctx.deps.history_processors)
+        message_history = await _process_message_history(
+            ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
+        )
         model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
         ctx.state.usage.incr(_usage.Usage())
 
@@ -762,7 +779,12 @@ async def _tool_from_mcp_server(
         # 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)
+
+        if server.process_tool_call is not None:
+            result = await server.process_tool_call(ctx, server.call_tool, tool_name, args)
+        else:
+            result = await server.call_tool(tool_name, args)
+
         return result
 
     for server in ctx.deps.mcp_servers:
@@ -876,12 +898,24 @@ def build_agent_graph(
 
 async def _process_message_history(
     messages: list[_messages.ModelMessage],
-    processors: Sequence[HistoryProcessor],
+    processors: Sequence[HistoryProcessor[DepsT]],
+    run_context: RunContext[DepsT],
 ) -> list[_messages.ModelMessage]:
     """Process message history through a sequence of processors."""
     for processor in processors:
+        takes_ctx = is_takes_ctx(processor)
+
         if is_async_callable(processor):
-            messages = await processor(messages)
+            if takes_ctx:
+                messages = await processor(run_context, messages)
+            else:
+                async_processor = cast(_HistoryProcessorAsync, processor)
+                messages = await async_processor(messages)
         else:
-            messages = await run_in_executor(processor, messages)
+            if takes_ctx:
+                sync_processor_with_ctx = cast(_HistoryProcessorSyncWithCtx[DepsT], processor)
+                messages = await run_in_executor(sync_processor_with_ctx, run_context, messages)
+            else:
+                sync_processor = cast(_HistoryProcessorSync, processor)
+                messages = await run_in_executor(sync_processor, messages)
     return messages

{pydantic_ai_slim-0.2.19 → pydantic_ai_slim-0.2.20}/pydantic_ai/_function_schema.py

@@ -8,7 +8,7 @@ from __future__ import annotations as _annotations
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, cast
+from typing import TYPE_CHECKING, Any, Callable, Union, cast
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -17,7 +17,7 @@ from pydantic.fields import FieldInfo
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import get_origin
+from typing_extensions import Concatenate, ParamSpec, TypeIs, TypeVar, get_origin
 
 from pydantic_ai.tools import RunContext
 
@@ -218,7 +218,16 @@ def function_schema(  # noqa: C901
     )
 
 
-def _takes_ctx(function: Callable[..., Any]) -> bool:
+P = ParamSpec('P')
+R = TypeVar('R')
+
+
+WithCtx = Callable[Concatenate[RunContext[Any], P], R]
+WithoutCtx = Callable[P, R]
+TargetFunc = Union[WithCtx[P, R], WithoutCtx[P, R]]
+
+
+def _takes_ctx(function: TargetFunc[P, R]) -> TypeIs[WithCtx[P, R]]:
     """Check if a function takes a `RunContext` first argument.
 
     Args:

{pydantic_ai_slim-0.2.19 → pydantic_ai_slim-0.2.20}/pydantic_ai/agent.py

@@ -180,7 +180,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
-        history_processors: Sequence[HistoryProcessor] | None = None,
+        history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
     ) -> None: ...
 
     @overload
@@ -210,7 +210,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
-        history_processors: Sequence[HistoryProcessor] | None = None,
+        history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
     ) -> None: ...
 
     def __init__(
@@ -235,7 +235,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
-        history_processors: Sequence[HistoryProcessor] | None = None,
+        history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
         **_deprecated_kwargs: Any,
     ):
         """Create an agent.

{pydantic_ai_slim-0.2.19 → pydantic_ai_slim-0.2.20}/pydantic_ai/mcp.py

@@ -4,7 +4,7 @@ import base64
 import functools
 import json
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Sequence
+from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
@@ -15,14 +15,20 @@ import anyio
 import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
+from mcp.shared.exceptions import McpError
 from mcp.shared.message import SessionMessage
 from mcp.types import (
     AudioContent,
     BlobResourceContents,
+    CallToolRequest,
+    CallToolRequestParams,
+    CallToolResult,
+    ClientRequest,
     Content,
     EmbeddedResource,
     ImageContent,
     LoggingLevel,
+    RequestParams,
     TextContent,
     TextResourceContents,
 )
@@ -30,7 +36,7 @@ from typing_extensions import Self, assert_never, deprecated
 
 from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.messages import BinaryContent
-from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.tools import RunContext, ToolDefinition
 
 try:
     from mcp.client.session import ClientSession
@@ -60,6 +66,9 @@ class MCPServer(ABC):
     e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
+    process_tool_call: ProcessToolCallback | None = None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     _client: ClientSession
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
@@ -113,13 +122,17 @@ class MCPServer(ABC):
         ]
 
     async def call_tool(
-        self, tool_name: str, arguments: dict[str, Any]
-    ) -> str | BinaryContent | dict[str, Any] | list[Any] | Sequence[str | BinaryContent | dict[str, Any] | list[Any]]:
+        self,
+        tool_name: str,
+        arguments: 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.
+            metadata: Request-level metadata (optional)
 
         Returns:
             The result of the tool call.
@@ -127,7 +140,23 @@ class MCPServer(ABC):
         Raises:
             ModelRetry: If the tool call fails.
         """
-        result = await self._client.call_tool(self.get_unprefixed_tool_name(tool_name), arguments)
+        try:
+            # meta param is not provided by session yet, so build and can send_request directly.
+            result = await self._client.send_request(
+                ClientRequest(
+                    CallToolRequest(
+                        method='tools/call',
+                        params=CallToolRequestParams(
+                            name=self.get_unprefixed_tool_name(tool_name),
+                            arguments=arguments,
+                            _meta=RequestParams.Meta(**metadata) if metadata else None,
+                        ),
+                    )
+                ),
+                CallToolResult,
+            )
+        except McpError as e:
+            raise ModelRetry(e.error.message)
 
         content = [self._map_tool_result_part(part) for part in result.content]
 
@@ -265,6 +294,9 @@ class MCPServerStdio(MCPServer):
     e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
+    process_tool_call: ProcessToolCallback | None = None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     timeout: float = 5
     """ The timeout in seconds to wait for the client to initialize."""
 
@@ -359,6 +391,9 @@ class _MCPServerHTTP(MCPServer):
     For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
+    process_tool_call: ProcessToolCallback | None = None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     @property
     @abstractmethod
     def _transport_client(
@@ -517,3 +552,29 @@ class MCPServerStreamableHTTP(_MCPServerHTTP):
     @property
     def _transport_client(self):
         return streamablehttp_client  # pragma: no cover
+
+
+ToolResult = (
+    str | BinaryContent | dict[str, Any] | list[Any] | Sequence[str | BinaryContent | dict[str, Any] | list[Any]]
+)
+"""The result type of a tool call."""
+
+CallToolFunc = Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[ToolResult]]
+"""A function type that represents a tool call."""
+
+ProcessToolCallback = Callable[
+    [
+        RunContext[Any],
+        CallToolFunc,
+        str,
+        dict[str, Any],
+    ],
+    Awaitable[ToolResult],
+]
+"""A process tool callback.
+
+It accepts a run context, the original tool call function, a tool name, and arguments.
+
+Allows wrapping an MCP server tool call to customize it, including adding extra request
+metadata.
+"""

{pydantic_ai_slim-0.2.19 → pydantic_ai_slim-0.2.20}/pydantic_ai/models/google.py

@@ -10,9 +10,8 @@ from uuid import uuid4
 
 from typing_extensions import assert_never
 
-from pydantic_ai.providers import Provider
-
 from .. import UnexpectedModelBehavior, _utils, usage
+from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
     FileUrl,
@@ -30,6 +29,7 @@ from ..messages import (
     VideoUrl,
 )
 from ..profiles import ModelProfileSpec
+from ..providers import Provider
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
@@ -52,6 +52,7 @@ try:
         FunctionDeclarationDict,
         GenerateContentConfigDict,
         GenerateContentResponse,
+        HttpOptionsDict,
         Part,
         PartDict,
         SafetySettingDict,
@@ -252,8 +253,17 @@ class GoogleModel(Model):
         tool_config = self._get_tool_config(model_request_parameters, tools)
         system_instruction, contents = await self._map_messages(messages)
 
+        http_options: HttpOptionsDict = {
+            'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
+        }
+        if timeout := model_settings.get('timeout'):
+            if isinstance(timeout, (int, float)):
+                http_options['timeout'] = int(1000 * timeout)
+            else:
+                raise UserError('Google does not support setting ModelSettings.timeout to a httpx.Timeout')
+
         config = GenerateContentConfigDict(
-            http_options={'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}},
+            http_options=http_options,
             system_instruction=system_instruction,
             temperature=model_settings.get('temperature'),
             top_p=model_settings.get('top_p'),