Improve test coverage

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -538,16 +538,17 @@ async def _handle_text_response(
 
         text = '\n\n'.join(texts)
         try:
+            run_context = build_run_context(ctx)
             if isinstance(output_schema, _output.TextOutputSchema):
-                run_context = build_run_context(ctx)
                 result_data = await output_schema.process(text, run_context)
             else:
                 m = _messages.RetryPromptPart(
                     content='Plain text responses are not permitted, please include your response in a tool call',
                 )
                 raise ToolRetryError(m)
 
-            result_data = await _validate_output(result_data, ctx, None)
+            for validator in ctx.deps.output_validators:
+                result_data = await validator.validate(result_data, run_context)
         except ToolRetryError as e:
             ctx.state.increment_retries(ctx.deps.max_result_retries, e)
             return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
@@ -834,17 +835,6 @@ async def _call_tool(
     return await toolset.call_tool(run_context, tool_call.tool_name, args_dict)
 
 
-async def _validate_output(
-    result_data: T,
-    ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, T]],
-    tool_call: _messages.ToolCallPart | None,
-) -> T:
-    for validator in ctx.deps.output_validators:
-        run_context = build_run_context(ctx)
-        result_data = await validator.validate(result_data, tool_call, run_context)
-    return result_data
-
-
 @dataclasses.dataclass
 class _RunMessages:
     messages: list[_messages.ModelMessage]

pydantic_ai_slim/pydantic_ai/_output.py

@@ -4,7 +4,7 @@
 import json
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Sequence
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast, overload
 
 from pydantic import TypeAdapter, ValidationError
@@ -83,28 +83,21 @@ def __post_init__(self):
     async def validate(
         self,
         result: T,
-        tool_call: _messages.ToolCallPart | None,
         run_context: RunContext[AgentDepsT],
         wrap_validation_errors: bool = True,
     ) -> T:
         """Validate a result but calling the function.
 
         Args:
             result: The result data after Pydantic validation the message content.
-            tool_call: The original tool call message, `None` if there was no tool call.
             run_context: The current run context.
             wrap_validation_errors: If true, wrap the validation errors in a retry message.
 
         Returns:
             Result of either the validated result data (ok) or a retry message (Err).
         """
         if self._takes_ctx:
-            ctx = (
-                replace(run_context, tool_name=tool_call.tool_name, tool_call_id=tool_call.tool_call_id)
-                if tool_call
-                else run_context
-            )
-            args = ctx, result
+            args = run_context, result
         else:
             args = (result,)
 
@@ -117,10 +110,12 @@ async def validate(
                 result_data = await _utils.run_in_executor(function, *args)
         except ModelRetry as r:
             if wrap_validation_errors:
-                m = _messages.RetryPromptPart(content=r.message)
-                if tool_call is not None:
-                    m.tool_name = tool_call.tool_name
-                    m.tool_call_id = tool_call.tool_call_id
+                m = _messages.RetryPromptPart(
+                    content=r.message,
+                    tool_name=run_context.tool_name,
+                )
+                if run_context.tool_call_id:
+                    m.tool_call_id = run_context.tool_call_id
                 raise ToolRetryError(m) from r
             else:
                 raise r
@@ -190,7 +185,7 @@ def build(  # noqa: C901
 
         if output := next((output for output in outputs if isinstance(output, NativeOutput)), None):
             if len(outputs) > 1:
-                raise UserError('`NativeOutput` must be the only output type.')
+                raise UserError('`NativeOutput` must be the only output type.')  # pragma: no cover
 
             return NativeOutputSchema(
                 processor=cls._build_processor(
@@ -203,7 +198,7 @@ def build(  # noqa: C901
             )
         elif output := next((output for output in outputs if isinstance(output, PromptedOutput)), None):
             if len(outputs) > 1:
-                raise UserError('`PromptedOutput` must be the only output type.')
+                raise UserError('`PromptedOutput` must be the only output type.')  # pragma: no cover
 
             return PromptedOutputSchema(
                 processor=cls._build_processor(
@@ -940,7 +935,7 @@ async def call_tool(
     ) -> Any:
         output = await self.processors[name].call(tool_args, ctx)
         for validator in self.output_validators:
-            output = await validator.validate(output, None, ctx, wrap_validation_errors=False)
+            output = await validator.validate(output, ctx, wrap_validation_errors=False)
         return output
 
 
@@ -965,7 +960,7 @@ def _flatten_output_spec(output_spec: OutputSpec[T]) -> Sequence[_OutputSpecItem
     for output in outputs:
         if isinstance(output, Sequence):
             outputs_flat.extend(_flatten_output_spec(cast(OutputSpec[T], output)))
-        if union_types := _utils.get_union_args(output):
+        elif union_types := _utils.get_union_args(output):
             outputs_flat.extend(union_types)
         else:
             outputs_flat.append(cast(_OutputSpecItem[T], output))

pydantic_ai_slim/pydantic_ai/agent.py

@@ -370,11 +370,10 @@ def __init__(
             output_retries = result_retries
 
         if mcp_servers := _deprecated_kwargs.pop('mcp_servers', None):
+            if toolsets is not None:  # pragma: no cover
+                raise TypeError('`mcp_servers` and `toolsets` cannot be set at the same time.')
             warnings.warn('`mcp_servers` is deprecated, use `toolsets` instead', DeprecationWarning)
-            if toolsets is None:
-                toolsets = mcp_servers
-            else:
-                toolsets = [*toolsets, *mcp_servers]
+            toolsets = mcp_servers
 
         _utils.validate_empty_kwargs(_deprecated_kwargs)
 

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -59,6 +59,7 @@ class MCPServer(AbstractToolset[Any], ABC):
     process_tool_call: ToolProcessFunc[Any] | None = None
     allow_sampling: bool = True
     max_retries: int = 1
+    sampling_model: models.Model | None = None
     # } end of "abstract fields"
 
     _running_count: int = 0
@@ -67,7 +68,6 @@ class MCPServer(AbstractToolset[Any], ABC):
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
     _exit_stack: AsyncExitStack
-    sampling_model: models.Model | None = None
 
     @abstractmethod
     @asynccontextmanager
@@ -83,11 +83,6 @@ async def client_streams(
         raise NotImplementedError('MCP Server subclasses must implement this method.')
         yield
 
-    @property
-    def is_running(self) -> bool:
-        """Check if the MCP server is running."""
-        return bool(self._running_count)
-
     @property
     def name(self) -> str:
         return repr(self)
@@ -373,6 +368,9 @@ async def main():
     max_retries: int = 1
     """The maximum number of times to retry a tool call."""
 
+    sampling_model: models.Model | None = None
+    """The model to use for sampling."""
+
     @asynccontextmanager
     async def client_streams(
         self,
@@ -471,6 +469,9 @@ class _MCPServerHTTP(MCPServer):
     max_retries: int = 1
     """The maximum number of times to retry a tool call."""
 
+    sampling_model: models.Model | None = None
+    """The model to use for sampling."""
+
     @property
     @abstractmethod
     def _transport_client(

pydantic_ai_slim/pydantic_ai/result.py

@@ -113,8 +113,8 @@ async def _validate_response(
             return await self._toolset.call_tool(run_context, tool_call.tool_name, args_dict)
         elif deferred_tool_calls := self._toolset.get_deferred_tool_calls(message.parts):
             if not self._output_schema.allows_deferred_tool_calls:
-                raise exceptions.UserError(
-                    'There are deferred tool calls but DeferredToolCalls is not among output types.'
+                raise exceptions.UserError(  # pragma: no cover
+                    'A deferred tool call was present, but `DeferredToolCalls` is not among output types. To resolve this, add `DeferredToolCalls` to the list of output types for this agent.'
                 )
             return cast(OutputDataT, deferred_tool_calls)
         elif isinstance(self._output_schema, TextOutputSchema):
@@ -124,7 +124,7 @@ async def _validate_response(
                 text, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )
             for validator in self._output_validators:
-                result_data = await validator.validate(result_data, None, self._run_ctx)
+                result_data = await validator.validate(result_data, self._run_ctx)
             return result_data
         else:
             raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
@@ -450,7 +450,7 @@ async def validate_structured_output(
         elif deferred_tool_calls := self._toolset.get_deferred_tool_calls(message.parts):
             if not self._output_schema.allows_deferred_tool_calls:
                 raise exceptions.UserError(
-                    'There are deferred tool calls but DeferredToolCalls is not among output types.'
+                    'A deferred tool call was present, but `DeferredToolCalls` is not among output types. To resolve this, add `DeferredToolCalls` to the list of output types for this agent.'
                 )
             return cast(OutputDataT, deferred_tool_calls)
         elif isinstance(self._output_schema, TextOutputSchema):
@@ -460,7 +460,7 @@ async def validate_structured_output(
                 text, self._run_ctx, allow_partial=allow_partial, wrap_validation_errors=False
             )
             for validator in self._output_validators:
-                result_data = await validator.validate(result_data, None, self._run_ctx)  # pragma: no cover
+                result_data = await validator.validate(result_data, self._run_ctx)  # pragma: no cover
             return result_data
         else:
             raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
@@ -469,7 +469,7 @@ async def validate_structured_output(
 
     async def _validate_text_output(self, text: str) -> str:
         for validator in self._output_validators:
-            text = await validator.validate(text, None, self._run_ctx)  # pragma: no cover
+            text = await validator.validate(text, self._run_ctx)  # pragma: no cover
         return text
 
     async def _marked_completed(self, message: _messages.ModelResponse) -> None:

pydantic_ai_slim/pydantic_ai/toolsets/_run.py

@@ -122,21 +122,22 @@ def _with_retry(self, name: str, ctx: RunContext[AgentDepsT]) -> Iterator[RunCon
             if current_retry == max_retries:
                 raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
             else:
-                if ctx.tool_call_id:
-                    if isinstance(e, ValidationError):
-                        m = _messages.RetryPromptPart(
-                            tool_name=name,
-                            content=e.errors(include_url=False, include_context=False),
-                            tool_call_id=ctx.tool_call_id,
-                        )
-                        e = ToolRetryError(m)
-                    elif isinstance(e, ModelRetry):
-                        m = _messages.RetryPromptPart(
-                            tool_name=name,
-                            content=e.message,
-                            tool_call_id=ctx.tool_call_id,
-                        )
-                        e = ToolRetryError(m)
+                if isinstance(e, ValidationError):
+                    m = _messages.RetryPromptPart(
+                        tool_name=name,
+                        content=e.errors(include_url=False, include_context=False),
+                    )
+                    if ctx.tool_call_id:  # pragma: no branch
+                        m.tool_call_id = ctx.tool_call_id
+                    e = ToolRetryError(m)
+                elif isinstance(e, ModelRetry):
+                    m = _messages.RetryPromptPart(
+                        tool_name=name,
+                        content=e.message,
+                    )
+                    if ctx.tool_call_id:  # pragma: no branch
+                        m.tool_call_id = ctx.tool_call_id
+                    e = ToolRetryError(m)
 
                 self._retries[name] = current_retry + 1
                 raise e

tests/test_agent.py

@@ -24,6 +24,7 @@
     ToolOutputSchema,
 )
 from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.mcp import MCPServerStdio
 from pydantic_ai.messages import (
     BinaryContent,
     ImageUrl,
@@ -46,7 +47,9 @@
 from pydantic_ai.profiles import ModelProfile
 from pydantic_ai.result import Usage
 from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.toolsets.combined import CombinedToolset
 from pydantic_ai.toolsets.function import FunctionToolset
+from pydantic_ai.toolsets.prefixed import PrefixedToolset
 
 from .conftest import IsDatetime, IsNow, IsStr, TestEnv
 
@@ -3439,6 +3442,14 @@ def test_deprecated_kwargs_still_work():
         assert issubclass(w[0].category, DeprecationWarning)
         assert '`result_retries` is deprecated' in str(w[0].message)
 
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter('always')
+
+        agent = Agent('test', mcp_servers=[MCPServerStdio('python', ['-m', 'tests.mcp_server'])])  # type: ignore[call-arg]
+        assert len(w) == 1
+        assert issubclass(w[0].category, DeprecationWarning)
+        assert '`mcp_servers` is deprecated' in str(w[0].message)
+
 
 def test_deprecated_kwargs_mixed_valid_invalid():
     """Test that mix of valid deprecated and invalid kwargs raises error for invalid ones."""
@@ -3583,3 +3594,39 @@ async def only_if_plan_presented(
             ),
         ]
     )
+
+
+async def test_reentrant_context_manager():
+    agent = Agent('test')
+    async with agent:
+        async with agent:
+            pass
+
+
+def test_set_mcp_sampling_model():
+    test_model = TestModel()
+    server1 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    server2 = MCPServerStdio('python', ['-m', 'tests.mcp_server'], sampling_model=test_model)
+    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix_')])
+    agent = Agent(None, toolsets=[toolset])
+
+    with pytest.raises(UserError, match='No sampling model provided and no model set on the agent.'):
+        agent.set_mcp_sampling_model()
+    assert server1.sampling_model is None
+    assert server2.sampling_model is test_model
+
+    agent.model = test_model
+    agent.set_mcp_sampling_model()
+    assert server1.sampling_model is test_model
+    assert server2.sampling_model is test_model
+
+    function_model = FunctionModel(lambda messages, info: ModelResponse(parts=[TextPart('Hello')]))
+    with agent.override(model=function_model):
+        agent.set_mcp_sampling_model()
+        assert server1.sampling_model is function_model
+        assert server2.sampling_model is function_model
+
+    function_model2 = FunctionModel(lambda messages, info: ModelResponse(parts=[TextPart('Goodbye')]))
+    agent.set_mcp_sampling_model(function_model2)
+    assert server1.sampling_model is function_model2
+    assert server2.sampling_model is function_model2

tests/test_examples.py

@@ -262,8 +262,6 @@ def rich_prompt_ask(prompt: str, *_args: Any, **_kwargs: Any) -> str:
 
 
 class MockMCPServer(AbstractToolset[Any]):
-    is_running = True
-
     async def __aenter__(self) -> MockMCPServer:
         return self