Remove Agent sampling_model field (and method argument) in favor of Agent.set_mcp_sampling_model

Author: DouweM

docs/mcp/client.md

@@ -364,6 +364,7 @@ agent = Agent('openai:gpt-4o', toolsets=[server])
 
 async def main():
     async with agent:
+        agent.set_mcp_sampling_model()
         result = await agent.run('Create an image of a robot in a punk style.')
     print(result.output)
     #> Image file written to robot_punk.svg.

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -115,7 +115,6 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     history_processors: Sequence[HistoryProcessor[DepsT]]
 
     toolset: RunToolset[DepsT]
-    sampling_model: models.Model
 
     tracer: Tracer
     instrumentation_settings: InstrumentationSettings | None = None
@@ -562,7 +561,6 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
         deps=ctx.deps.user_deps,
         model=ctx.deps.model,
         usage=ctx.state.usage,
-        sampling_model=ctx.deps.sampling_model,
         prompt=ctx.deps.prompt,
         messages=ctx.state.message_history,
         run_step=ctx.state.run_step,

pydantic_ai_slim/pydantic_ai/_run_context.py

@@ -27,8 +27,6 @@ class RunContext(Generic[AgentDepsT]):
     """The model used in this run."""
     usage: Usage
     """LLM usage associated with the run."""
-    sampling_model: Model
-    """The model used for MCP sampling."""
     prompt: str | Sequence[_messages.UserContent] | None = None
     """The original user prompt passed to the run."""
     messages: list[_messages.ModelMessage] = field(default_factory=list)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -3,9 +3,8 @@
 import base64
 import functools
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Iterator, Sequence
-from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager, nullcontext
-from contextvars import ContextVar
+from collections.abc import AsyncIterator, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
 from types import TracebackType
@@ -70,22 +69,6 @@ class MCPServer(AbstractToolset[Any], ABC):
     _exit_stack: AsyncExitStack
     sampling_model: models.Model | None = None
 
-    def __post_init__(self):
-        self._override_sampling_model: ContextVar[models.Model | None] = ContextVar(
-            '_override_sampling_model', default=None
-        )
-
-    @contextmanager
-    def override_sampling_model(
-        self,
-        model: models.Model,
-    ) -> Iterator[None]:
-        token = self._override_sampling_model.set(model)
-        try:
-            yield
-        finally:
-            self._override_sampling_model.reset(token)
-
     @abstractmethod
     @asynccontextmanager
     async def client_streams(
@@ -149,28 +132,23 @@ async def call_tool(
         Raises:
             ModelRetry: If the tool call fails.
         """
-        sampling_contextmanager = (
-            nullcontext() if self._get_sampling_model() else self.override_sampling_model(ctx.sampling_model)
-        )
-        with sampling_contextmanager:
-            async with self:  # Ensure server is running
-                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=name,
-                                    arguments=tool_args,
-                                    _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=tool_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]
 
@@ -245,15 +223,11 @@ async def __aexit__(
         if self._running_count <= 0:
             await self._exit_stack.aclose()
 
-    def _get_sampling_model(self) -> models.Model | None:
-        return self._override_sampling_model.get() or self.sampling_model
-
     async def _sampling_callback(
         self, context: RequestContext[ClientSession, Any], params: mcp_types.CreateMessageRequestParams
     ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData:
         """MCP sampling callback."""
-        sampling_model = self._get_sampling_model()
-        if sampling_model is None:
+        if self.sampling_model is None:
             raise ValueError('Sampling model is not set')  # pragma: no cover
 
         pai_messages = _mcp.map_from_mcp_params(params)
@@ -265,15 +239,15 @@ async def _sampling_callback(
         if stop_sequences := params.stopSequences:  # pragma: no branch
             model_settings['stop_sequences'] = stop_sequences
 
-        model_response = await sampling_model.request(
+        model_response = await self.sampling_model.request(
             pai_messages,
             model_settings,
             models.ModelRequestParameters(),
         )
         return mcp_types.CreateMessageResult(
             role='assistant',
             content=_mcp.map_from_model_response(model_response),
-            model=sampling_model.model_name,
+            model=self.sampling_model.model_name,
         )
 
     def _map_tool_result_part(

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Generic, Literal
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal
 
 from pydantic_core import SchemaValidator
 from typing_extensions import Self
@@ -78,3 +78,6 @@ async def call_tool(
         self, ctx: RunContext[AgentDepsT], name: str, tool_args: dict[str, Any], *args: Any, **kwargs: Any
     ) -> Any:
         raise NotImplementedError()
+
+    def accept(self, visitor: Callable[[AbstractToolset[AgentDepsT]], Any]) -> Any:
+        return visitor(self)

pydantic_ai_slim/pydantic_ai/toolsets/__init__.py

@@ -5,7 +5,7 @@
 from contextlib import AsyncExitStack
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Any
+from typing import Any, Callable
 
 from pydantic_core import SchemaValidator
 from typing_extensions import Self
@@ -89,6 +89,10 @@ async def call_tool(
     ) -> Any:
         return await self._toolset_for_tool_name(name).call_tool(ctx, name, tool_args, *args, **kwargs)
 
+    def accept(self, visitor: Callable[[AbstractToolset[AgentDepsT]], Any]) -> Any:
+        for toolset in self.toolsets:
+            toolset.accept(visitor)
+
     def _toolset_for_tool_name(self, name: str) -> AbstractToolset[AgentDepsT]:
         try:
             return self._toolset_per_tool_name[name]

pydantic_ai_slim/pydantic_ai/toolsets/combined.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from types import TracebackType
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
 from pydantic_core import SchemaValidator
 from typing_extensions import Self
@@ -58,5 +58,8 @@ async def call_tool(
     ) -> Any:
         return await self.wrapped.call_tool(ctx, name, tool_args, *args, **kwargs)
 
+    def accept(self, visitor: Callable[[AbstractToolset[AgentDepsT]], Any]) -> Any:
+        return self.wrapped.accept(visitor)
+
     def __getattr__(self, item: str):
         return getattr(self.wrapped, item)  # pragma: no cover

pydantic_ai_slim/pydantic_ai/toolsets/wrapper.py

@@ -6,7 +6,6 @@
 import shutil
 import sys
 from collections.abc import AsyncIterator, Iterable, Sequence
-from contextlib import nullcontext
 from dataclasses import dataclass
 from inspect import FrameInfo
 from io import StringIO
@@ -264,7 +263,6 @@ def rich_prompt_ask(prompt: str, *_args: Any, **_kwargs: Any) -> str:
 
 class MockMCPServer(AbstractToolset[Any]):
     is_running = True
-    override_sampling_model = nullcontext
 
     async def __aenter__(self) -> MockMCPServer:
         return self

tests/test_examples.py

@@ -68,7 +68,7 @@ def agent(model: Model, mcp_server: MCPServerStdio) -> Agent:
 
 @pytest.fixture
 def run_context(model: Model) -> RunContext[int]:
-    return RunContext(deps=0, model=model, usage=Usage(), sampling_model=model)
+    return RunContext(deps=0, model=model, usage=Usage())
 
 
 async def test_stdio_server(run_context: RunContext[int]):