From 8d4afa5fde719d334f7d56e5167f2878c45d8b9e Mon Sep 17 00:00:00 2001 From: Alex Hall Date: Tue, 2 Sep 2025 15:16:30 +0200 Subject: [PATCH] Minor fixes: Agent.__aenter__ returns Self, use default instrumentation for MCP sampling model --- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 4 ++-- pydantic_ai_slim/pydantic_ai/mcp.py | 7 ++----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 58df0d92d9..273c1ef944 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -12,7 +12,7 @@ import anyio from opentelemetry.trace import NoOpTracer, use_span from pydantic.json_schema import GenerateJsonSchema -from typing_extensions import TypeVar, deprecated +from typing_extensions import Self, TypeVar, deprecated from pydantic_graph import Graph @@ -1355,7 +1355,7 @@ def _prepare_output_schema( return schema # pyright: ignore[reportReturnType] - async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]: + async def __aenter__(self) -> Self: """Enter the agent context. This will start all [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] registered as `toolsets` so they are ready to be used. diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index f495d228b6..168405723b 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -20,6 +20,7 @@ from pydantic_ai.tools import RunContext, ToolDefinition +from .direct import model_request from .toolsets.abstract import AbstractToolset, ToolsetTool try: @@ -329,11 +330,7 @@ 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( - pai_messages, - model_settings, - models.ModelRequestParameters(), - ) + model_response = await model_request(self.sampling_model, pai_messages, model_settings=model_settings) return mcp_types.CreateMessageResult( role='assistant', content=_mcp.map_from_model_response(model_response),