Merge branch 'main' into toolsets

# Conflicts:
#	pydantic_ai_slim/pydantic_ai/agent.py
#	pydantic_ai_slim/pydantic_ai/mcp.py

Author: DouweM

pydantic_ai_slim/pydantic_ai/agent.py

@@ -5,7 +5,8 @@
 import json
 import warnings
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
-from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
+from contextvars import ContextVar
 from copy import deepcopy
 from types import FrameType
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, cast, final, overload
@@ -164,8 +165,6 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     _prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
-    _override_deps: _utils.Option[AgentDepsT] = dataclasses.field(default=None, repr=False)
-    _override_model: _utils.Option[models.Model] = dataclasses.field(default=None, repr=False)
 
     @overload
     def __init__(
@@ -425,6 +424,9 @@ def __init__(
 
         self.history_processors = history_processors or []
 
+        self._override_deps: ContextVar[_utils.Option[AgentDepsT]] = ContextVar('_override_deps', default=None)
+        self._override_model: ContextVar[_utils.Option[models.Model]] = ContextVar('_override_model', default=None)
+
     @staticmethod
     def instrument_all(instrument: InstrumentationSettings | bool = True) -> None:
         """Set the instrumentation options for all agents where `instrument` is not set."""
@@ -1212,24 +1214,22 @@ def override(
             model: The model to use instead of the model passed to the agent run.
         """
         if _utils.is_set(deps):
-            override_deps_before = self._override_deps
-            self._override_deps = _utils.Some(deps)
+            deps_token = self._override_deps.set(_utils.Some(deps))
         else:
-            override_deps_before = _utils.UNSET
+            deps_token = None
 
         if _utils.is_set(model):
-            override_model_before = self._override_model
-            self._override_model = _utils.Some(models.infer_model(model))
+            model_token = self._override_model.set(_utils.Some(models.infer_model(model)))
         else:
-            override_model_before = _utils.UNSET
+            model_token = None
 
         try:
             yield
         finally:
-            if _utils.is_set(override_deps_before):
-                self._override_deps = override_deps_before
-            if _utils.is_set(override_model_before):
-                self._override_model = override_model_before
+            if deps_token is not None:
+                self._override_deps.reset(deps_token)
+            if model_token is not None:
+                self._override_model.reset(model_token)
 
     @overload
     def instructions(
@@ -1662,7 +1662,7 @@ def _get_model(self, model: models.Model | models.KnownModelName | str | None) -
             The model used
         """
         model_: models.Model
-        if some_model := self._override_model:
+        if some_model := self._override_model.get():
             # we don't want `override()` to cover up errors from the model not being defined, hence this check
             if model is None and self.model is None:
                 raise exceptions.UserError(
@@ -1691,7 +1691,7 @@ def _get_deps(self: Agent[T, OutputDataT], deps: T) -> T:
 
         We could do runtime type checking of deps against `self._deps_type`, but that's a slippery slope.
         """
-        if some_deps := self._override_deps:
+        if some_deps := self._override_deps.get():
             return some_deps.value
         else:
             return deps
@@ -1793,10 +1793,11 @@ async def run_toolsets(
             model: models.Model | None = self._get_model(sampling_model)
         except exceptions.UserError:  # pragma: no cover
             model = None
-        if model is not None:  # pragma: no branch
-            self._toolset._set_mcp_sampling_model(model)  # type: ignore[reportPrivateUsage]
 
-        async with self._toolset:
+        async with AsyncExitStack() as exit_stack:
+            if model is not None:  # pragma: no branch
+                exit_stack.enter_context(self._toolset.override_sampling_model(model))
+            await exit_stack.enter_async_context(self._toolset)
             yield
 
     @asynccontextmanager

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -3,8 +3,9 @@
 import base64
 import functools
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Sequence
-from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
+from collections.abc import AsyncIterator, Iterator, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass
 from pathlib import Path
 from types import TracebackType
@@ -69,6 +70,22 @@ 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(
@@ -193,9 +210,6 @@ def _get_tool_args_validator(self, ctx: RunContext[Any], name: str) -> pydantic_
     def _max_retries_for_tool(self, name: str) -> int:
         return self.max_retries
 
-    def _set_mcp_sampling_model(self, model: models.Model) -> None:
-        self.sampling_model = model
-
     async def __aenter__(self) -> Self:
         if self._running_count == 0:
             self._exit_stack = AsyncExitStack()
@@ -231,7 +245,8 @@ async def _sampling_callback(
         self, context: RequestContext[ClientSession, Any], params: mcp_types.CreateMessageRequestParams
     ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData:
         """MCP sampling callback."""
-        if self.sampling_model is None:
+        sampling_model = self._override_sampling_model.get() or self.sampling_model
+        if sampling_model is None:
             raise ValueError('Sampling model is not set')  # pragma: no cover
 
         pai_messages = _mcp.map_from_mcp_params(params)
@@ -243,15 +258,15 @@ async def _sampling_callback(
         if stop_sequences := params.stopSequences:  # pragma: no branch
             model_settings['stop_sequences'] = stop_sequences
 
-        model_response = await self.sampling_model.request(
+        model_response = await sampling_model.request(
             pai_messages,
             model_settings,
             models.ModelRequestParameters(),
         )
         return mcp_types.CreateMessageResult(
             role='assistant',
             content=_mcp.map_from_model_response(model_response),
-            model=self.sampling_model.model_name,
+            model=sampling_model.model_name,
         )
 
     def _map_tool_result_part(

pydantic_ai_slim/pydantic_ai/toolsets/__init__.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from contextlib import contextmanager
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Generic, Literal
 
@@ -80,5 +82,6 @@ async def call_tool(
     ) -> Any:
         raise NotImplementedError()
 
-    def _set_mcp_sampling_model(self, model: Model) -> None:
-        pass
+    @contextmanager
+    def override_sampling_model(self, model: Model) -> Iterator[None]:
+        yield

pydantic_ai_slim/pydantic_ai/toolsets/combined.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Sequence
-from contextlib import AsyncExitStack
+from collections.abc import Iterator, Sequence
+from contextlib import AsyncExitStack, ExitStack, contextmanager
 from dataclasses import dataclass
 from types import TracebackType
 from typing import TYPE_CHECKING, Any
@@ -92,9 +92,12 @@ async def call_tool(
     ) -> Any:
         return await self._toolset_for_tool_name(name).call_tool(ctx, name, tool_args, *args, **kwargs)
 
-    def _set_mcp_sampling_model(self, model: Model) -> None:
-        for toolset in self.toolsets:
-            toolset._set_mcp_sampling_model(model)
+    @contextmanager
+    def override_sampling_model(self, model: Model) -> Iterator[None]:
+        with ExitStack() as exit_stack:
+            for toolset in self.toolsets:
+                exit_stack.enter_context(toolset.override_sampling_model(model))
+            yield
 
     def _toolset_for_tool_name(self, name: str) -> AbstractToolset[AgentDepsT]:
         try:

pydantic_ai_slim/pydantic_ai/toolsets/wrapper.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass
 from types import TracebackType
 from typing import TYPE_CHECKING, Any
@@ -59,8 +61,10 @@ async def call_tool(
     ) -> Any:
         return await self.wrapped.call_tool(ctx, name, tool_args, *args, **kwargs)
 
-    def _set_mcp_sampling_model(self, model: Model) -> None:
-        self.wrapped._set_mcp_sampling_model(model)
+    @contextmanager
+    def override_sampling_model(self, model: Model) -> Iterator[None]:
+        with self.wrapped.override_sampling_model(model):
+            yield
 
     def __getattr__(self, item: str):
         return getattr(self.wrapped, item)  # pragma: no cover

tests/test_examples.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