pydantic-ai-slim 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

pydantic_ai/_output.py

@@ -1,22 +1,55 @@
 from __future__ import annotations as _annotations
 
 import inspect
-from collections.abc import Awaitable, Iterable, Iterator
+from collections.abc import Awaitable, Iterable, Iterator, Sequence
 from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, Literal, Union, cast
 
 from pydantic import TypeAdapter, ValidationError
-from typing_extensions import TypedDict, TypeVar, get_args, get_origin
+from pydantic_core import SchemaValidator
+from typing_extensions import TypeAliasType, TypedDict, TypeVar, get_args, get_origin
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
-from . import _utils, messages as _messages
+from . import _function_schema, _utils, messages as _messages
 from .exceptions import ModelRetry
-from .result import DEFAULT_OUTPUT_TOOL_NAME, OutputDataT, OutputDataT_inv, OutputValidatorFunc, ToolOutput
-from .tools import AgentDepsT, GenerateToolJsonSchema, RunContext, ToolDefinition
+from .tools import AgentDepsT, GenerateToolJsonSchema, ObjectJsonSchema, RunContext, ToolDefinition
 
 T = TypeVar('T')
 """An invariant TypeVar."""
+OutputDataT_inv = TypeVar('OutputDataT_inv', default=str)
+"""
+An invariant type variable for the result data of a model.
+
+We need to use an invariant typevar for `OutputValidator` and `OutputValidatorFunc` because the output data type is used
+in both the input and output of a `OutputValidatorFunc`. This can theoretically lead to some issues assuming that types
+possessing OutputValidator's are covariant in the result data type, but in practice this is rarely an issue, and
+changing it would have negative consequences for the ergonomics of the library.
+
+At some point, it may make sense to change the input to OutputValidatorFunc to be `Any` or `object` as doing that would
+resolve these potential variance issues.
+"""
+OutputDataT = TypeVar('OutputDataT', default=str, covariant=True)
+"""Covariant type variable for the result data type of a run."""
+
+OutputValidatorFunc = Union[
+    Callable[[RunContext[AgentDepsT], OutputDataT_inv], OutputDataT_inv],
+    Callable[[RunContext[AgentDepsT], OutputDataT_inv], Awaitable[OutputDataT_inv]],
+    Callable[[OutputDataT_inv], OutputDataT_inv],
+    Callable[[OutputDataT_inv], Awaitable[OutputDataT_inv]],
+]
+"""
+A function that always takes and returns the same type of data (which is the result type of an agent run), and:
+
+* may or may not take [`RunContext`][pydantic_ai.tools.RunContext] as a first argument
+* may or may not be async
+
+Usage `OutputValidatorFunc[AgentDepsT, T]`.
+"""
+
+
+DEFAULT_OUTPUT_TOOL_NAME = 'final_result'
+DEFAULT_OUTPUT_TOOL_DESCRIPTION = 'The final response which ends this conversation'
 
 
 @dataclass
@@ -76,69 +109,135 @@ class ToolRetryError(Exception):
         super().__init__()
 
 
+@dataclass(init=False)
+class ToolOutput(Generic[OutputDataT]):
+    """Marker class to use tools for outputs, and customize the tool."""
+
+    output_type: SimpleOutputType[OutputDataT]
+    name: str | None
+    description: str | None
+    max_retries: int | None
+    strict: bool | None
+
+    def __init__(
+        self,
+        type_: SimpleOutputType[OutputDataT],
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        max_retries: int | None = None,
+        strict: bool | None = None,
+    ):
+        self.output_type = type_
+        self.name = name
+        self.description = description
+        self.max_retries = max_retries
+        self.strict = strict
+
+
+T_co = TypeVar('T_co', covariant=True)
+# output_type=Type or output_type=function or output_type=object.method
+SimpleOutputType = TypeAliasType(
+    'SimpleOutputType', Union[type[T_co], Callable[..., T_co], Callable[..., Awaitable[T_co]]], type_params=(T_co,)
+)
+# output_type=ToolOutput(<see above>) or <see above>
+SimpleOutputTypeOrMarker = TypeAliasType(
+    'SimpleOutputTypeOrMarker', Union[SimpleOutputType[T_co], ToolOutput[T_co]], type_params=(T_co,)
+)
+# output_type=<see above> or [<see above>, ...]
+OutputType = TypeAliasType(
+    'OutputType', Union[SimpleOutputTypeOrMarker[T_co], Sequence[SimpleOutputTypeOrMarker[T_co]]], type_params=(T_co,)
+)
+
+
 @dataclass
 class OutputSchema(Generic[OutputDataT]):
-    """Model the final response from an agent run.
+    """Model the final output from an agent run.
 
     Similar to `Tool` but for the final output of running an agent.
     """
 
-    tools: dict[str, OutputSchemaTool[OutputDataT]]
+    tools: dict[str, OutputTool[OutputDataT]]
     allow_text_output: bool
 
     @classmethod
     def build(
-        cls: type[OutputSchema[T]],
-        output_type: type[T] | ToolOutput[T],
+        cls: type[OutputSchema[OutputDataT]],
+        output_type: OutputType[OutputDataT],
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
-    ) -> OutputSchema[T] | None:
-        """Build an OutputSchema dataclass from a response type."""
+    ) -> OutputSchema[OutputDataT] | None:
+        """Build an OutputSchema dataclass from an output type."""
         if output_type is str:
             return None
 
-        if isinstance(output_type, ToolOutput):
-            # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
-            name = output_type.name
-            description = output_type.description
-            output_type_ = output_type.output_type
-            strict = output_type.strict
+        output_types: Sequence[SimpleOutputTypeOrMarker[OutputDataT]]
+        if isinstance(output_type, Sequence):
+            output_types = output_type
         else:
-            output_type_ = output_type
+            output_types = (output_type,)
 
-        if output_type_option := extract_str_from_union(output_type):
-            output_type_ = output_type_option.value
+        output_types_flat: list[SimpleOutputTypeOrMarker[OutputDataT]] = []
+        for output_type in output_types:
+            if union_types := get_union_args(output_type):
+                output_types_flat.extend(union_types)
+            else:
+                output_types_flat.append(output_type)
+
+        allow_text_output = False
+        if str in output_types_flat:
             allow_text_output = True
-        else:
-            allow_text_output = False
-
-        tools: dict[str, OutputSchemaTool[T]] = {}
-        if args := get_union_args(output_type_):
-            for i, arg in enumerate(args, start=1):
-                tool_name = raw_tool_name = union_tool_name(name, arg)
-                while tool_name in tools:
-                    tool_name = f'{raw_tool_name}_{i}'
-                tools[tool_name] = cast(
-                    OutputSchemaTool[T],
-                    OutputSchemaTool(
-                        output_type=arg, name=tool_name, description=description, multiple=True, strict=strict
-                    ),
-                )
-        else:
-            name = name or DEFAULT_OUTPUT_TOOL_NAME
-            tools[name] = cast(
-                OutputSchemaTool[T],
-                OutputSchemaTool(
-                    output_type=output_type_, name=name, description=description, multiple=False, strict=strict
-                ),
+            output_types_flat = [t for t in output_types_flat if t is not str]
+
+        multiple = len(output_types_flat) > 1
+
+        default_tool_name = name or DEFAULT_OUTPUT_TOOL_NAME
+        default_tool_description = description
+        default_tool_strict = strict
+
+        tools: dict[str, OutputTool[OutputDataT]] = {}
+        for output_type in output_types_flat:
+            tool_name = None
+            tool_description = None
+            tool_strict = None
+            if isinstance(output_type, ToolOutput):
+                tool_output_type = output_type.output_type
+                # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
+                tool_name = output_type.name
+                tool_description = output_type.description
+                tool_strict = output_type.strict
+            else:
+                tool_output_type = output_type
+
+            if tool_name is None:
+                tool_name = default_tool_name
+                if multiple:
+                    tool_name += f'_{tool_output_type.__name__}'
+
+            i = 1
+            original_tool_name = tool_name
+            while tool_name in tools:
+                i += 1
+                tool_name = f'{original_tool_name}_{i}'
+
+            tool_description = tool_description or default_tool_description
+            if tool_strict is None:
+                tool_strict = default_tool_strict
+
+            parameters_schema = OutputObjectSchema(
+                output_type=tool_output_type, description=tool_description, strict=tool_strict
             )
+            tools[tool_name] = OutputTool(name=tool_name, parameters_schema=parameters_schema, multiple=multiple)
 
-        return cls(tools=tools, allow_text_output=allow_text_output)
+        return cls(
+            tools=tools,
+            allow_text_output=allow_text_output,
+        )
 
     def find_named_tool(
         self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
-    ) -> tuple[_messages.ToolCallPart, OutputSchemaTool[OutputDataT]] | None:
+    ) -> tuple[_messages.ToolCallPart, OutputTool[OutputDataT]] | None:
         """Find a tool that matches one of the calls, with a specific name."""
         for part in parts:  # pragma: no branch
             if isinstance(part, _messages.ToolCallPart):  # pragma: no branch
@@ -148,7 +247,7 @@ class OutputSchema(Generic[OutputDataT]):
     def find_tool(
         self,
         parts: Iterable[_messages.ModelResponsePart],
-    ) -> Iterator[tuple[_messages.ToolCallPart, OutputSchemaTool[OutputDataT]]]:
+    ) -> Iterator[tuple[_messages.ToolCallPart, OutputTool[OutputDataT]]]:
         """Find a tool that matches one of the calls."""
         for part in parts:
             if isinstance(part, _messages.ToolCallPart):  # pragma: no branch
@@ -164,64 +263,138 @@ class OutputSchema(Generic[OutputDataT]):
         return [t.tool_def for t in self.tools.values()]
 
 
-DEFAULT_DESCRIPTION = 'The final response which ends this conversation'
+def allow_text_output(output_schema: OutputSchema[Any] | None) -> bool:
+    return output_schema is None or output_schema.allow_text_output
+
+
+@dataclass
+class OutputObjectDefinition:
+    name: str
+    json_schema: ObjectJsonSchema
+    description: str | None = None
+    strict: bool | None = None
 
 
 @dataclass(init=False)
-class OutputSchemaTool(Generic[OutputDataT]):
-    tool_def: ToolDefinition
-    type_adapter: TypeAdapter[Any]
+class OutputObjectSchema(Generic[OutputDataT]):
+    definition: OutputObjectDefinition
+    validator: SchemaValidator
+    function_schema: _function_schema.FunctionSchema | None = None
+    outer_typed_dict_key: str | None = None
 
     def __init__(
-        self, *, output_type: type[OutputDataT], name: str, description: str | None, multiple: bool, strict: bool | None
+        self,
+        *,
+        output_type: SimpleOutputType[OutputDataT],
+        name: str | None = None,
+        description: str | None = None,
+        strict: bool | None = None,
     ):
-        """Build a OutputSchemaTool from a response type."""
-        if _utils.is_model_like(output_type):
-            self.type_adapter = TypeAdapter(output_type)
-            outer_typed_dict_key: str | None = None
-            # noinspection PyArgumentList
-            parameters_json_schema = _utils.check_object_json_schema(
-                self.type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
-            )
+        if inspect.isfunction(output_type) or inspect.ismethod(output_type):
+            self.function_schema = _function_schema.function_schema(output_type, GenerateToolJsonSchema)
+            self.validator = self.function_schema.validator
+            json_schema = self.function_schema.json_schema
+            json_schema['description'] = self.function_schema.description
         else:
-            response_data_typed_dict = TypedDict(  # noqa: UP013
-                'response_data_typed_dict',
-                {'response': output_type},  # pyright: ignore[reportInvalidTypeForm]
-            )
-            self.type_adapter = TypeAdapter(response_data_typed_dict)
-            outer_typed_dict_key = 'response'
-            # noinspection PyArgumentList
-            parameters_json_schema = _utils.check_object_json_schema(
-                self.type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
+            type_adapter: TypeAdapter[Any]
+            if _utils.is_model_like(output_type):
+                type_adapter = TypeAdapter(output_type)
+            else:
+                self.outer_typed_dict_key = 'response'
+                response_data_typed_dict = TypedDict(  # noqa: UP013
+                    'response_data_typed_dict',
+                    {'response': cast(type[OutputDataT], output_type)},  # pyright: ignore[reportInvalidTypeForm]
+                )
+                type_adapter = TypeAdapter(response_data_typed_dict)
+
+            # Really a PluggableSchemaValidator, but it's API-compatible
+            self.validator = cast(SchemaValidator, type_adapter.validator)
+            json_schema = _utils.check_object_json_schema(
+                type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
             )
-            # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM
-            parameters_json_schema.pop('title')
 
-        if json_schema_description := parameters_json_schema.pop('description', None):
+            if self.outer_typed_dict_key:
+                # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM
+                json_schema.pop('title')
+
+        if json_schema_description := json_schema.pop('description', None):
             if description is None:
-                tool_description = json_schema_description
+                description = json_schema_description
             else:
-                tool_description = f'{description}. {json_schema_description}'  # pragma: no cover
+                description = f'{description}. {json_schema_description}'
+
+        self.definition = OutputObjectDefinition(
+            name=name or getattr(output_type, '__name__', DEFAULT_OUTPUT_TOOL_NAME),
+            description=description,
+            json_schema=json_schema,
+            strict=strict,
+        )
+
+    async def process(
+        self,
+        data: str | dict[str, Any] | None,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+    ) -> OutputDataT:
+        """Process an output message, performing validation and (if necessary) calling the output function.
+
+        Args:
+            data: The output data to validate.
+            run_context: The current run context.
+            allow_partial: If true, allow partial validation.
+
+        Returns:
+            Either the validated output data (left) or a retry message (right).
+        """
+        pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
+        if isinstance(data, str):
+            output = self.validator.validate_json(data or '{}', allow_partial=pyd_allow_partial)
         else:
-            tool_description = description or DEFAULT_DESCRIPTION
+            output = self.validator.validate_python(data or {}, allow_partial=pyd_allow_partial)
+
+        if self.function_schema:
+            output = await self.function_schema.call(output, run_context)
+
+        if k := self.outer_typed_dict_key:
+            output = output[k]
+        return output
+
+
+@dataclass(init=False)
+class OutputTool(Generic[OutputDataT]):
+    parameters_schema: OutputObjectSchema[OutputDataT]
+    tool_def: ToolDefinition
+
+    def __init__(self, *, name: str, parameters_schema: OutputObjectSchema[OutputDataT], multiple: bool):
+        self.parameters_schema = parameters_schema
+        definition = parameters_schema.definition
+
+        description = definition.description
+        if not description:
+            description = DEFAULT_OUTPUT_TOOL_DESCRIPTION
             if multiple:
-                tool_description = f'{union_arg_name(output_type)}: {tool_description}'
+                description = f'{definition.name}: {description}'
 
         self.tool_def = ToolDefinition(
             name=name,
-            description=tool_description,
-            parameters_json_schema=parameters_json_schema,
-            outer_typed_dict_key=outer_typed_dict_key,
-            strict=strict,
+            description=description,
+            parameters_json_schema=definition.json_schema,
+            strict=definition.strict,
+            outer_typed_dict_key=parameters_schema.outer_typed_dict_key,
         )
 
-    def validate(
-        self, tool_call: _messages.ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+    async def process(
+        self,
+        tool_call: _messages.ToolCallPart,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
     ) -> OutputDataT:
-        """Validate an output message.
+        """Process an output message.
 
         Args:
             tool_call: The tool call from the LLM to validate.
+            run_context: The current run context.
             allow_partial: If true, allow partial validation.
             wrap_validation_errors: If true, wrap the validation errors in a retry message.
 
@@ -229,57 +402,31 @@ class OutputSchemaTool(Generic[OutputDataT]):
             Either the validated output data (left) or a retry message (right).
         """
         try:
-            pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-            if isinstance(tool_call.args, str):
-                output = self.type_adapter.validate_json(
-                    tool_call.args or '{}', experimental_allow_partial=pyd_allow_partial
-                )
-            else:
-                output = self.type_adapter.validate_python(
-                    tool_call.args or {}, experimental_allow_partial=pyd_allow_partial
-                )
+            output = await self.parameters_schema.process(tool_call.args, run_context, allow_partial=allow_partial)
         except ValidationError as e:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(
                     tool_name=tool_call.tool_name,
-                    content=e.errors(include_url=False),
+                    content=e.errors(include_url=False, include_context=False),
                     tool_call_id=tool_call.tool_call_id,
                 )
                 raise ToolRetryError(m) from e
             else:
                 raise  # pragma: lax no cover
+        except ModelRetry as r:
+            if wrap_validation_errors:
+                m = _messages.RetryPromptPart(
+                    tool_name=tool_call.tool_name,
+                    content=r.message,
+                    tool_call_id=tool_call.tool_call_id,
+                )
+                raise ToolRetryError(m) from r
+            else:
+                raise  # pragma: lax no cover
         else:
-            if k := self.tool_def.outer_typed_dict_key:
-                output = output[k]
             return output
 
 
-def union_tool_name(base_name: str | None, union_arg: Any) -> str:
-    return f'{base_name or DEFAULT_OUTPUT_TOOL_NAME}_{union_arg_name(union_arg)}'
-
-
-def union_arg_name(union_arg: Any) -> str:
-    return union_arg.__name__
-
-
-def extract_str_from_union(output_type: Any) -> _utils.Option[Any]:
-    """Extract the string type from a Union, return the remaining union or remaining type."""
-    union_args = get_union_args(output_type)
-    if any(t is str for t in union_args):
-        remain_args: list[Any] = []
-        includes_str = False
-        for arg in union_args:
-            if arg is str:
-                includes_str = True
-            else:
-                remain_args.append(arg)
-        if includes_str:  # pragma: no branch
-            if len(remain_args) == 1:
-                return _utils.Some(remain_args[0])
-            else:
-                return _utils.Some(Union[tuple(remain_args)])  # pragma: no cover
-
-
 def get_union_args(tp: Any) -> tuple[Any, ...]:
     """Extract the arguments of a Union type if `output_type` is a union, otherwise return an empty tuple."""
     if typing_objects.is_typealiastype(tp):

pydantic_ai/agent.py

@@ -29,7 +29,7 @@ from . import (
     usage as _usage,
 )
 from .models.instrumented import InstrumentationSettings, InstrumentedModel, instrument_model
-from .result import FinalResult, OutputDataT, StreamedRunResult, ToolOutput
+from .result import FinalResult, OutputDataT, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
     AgentDepsT,
@@ -127,7 +127,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     be merged with this value, with the runtime argument taking priority.
     """
 
-    output_type: type[OutputDataT] | ToolOutput[OutputDataT]
+    output_type: _output.OutputType[OutputDataT]
     """
     The type of data output by agent runs, used to validate the data returned by the model, defaults to `str`.
     """
@@ -162,7 +162,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         model: models.Model | models.KnownModelName | str | None = None,
         *,
-        output_type: type[OutputDataT] | ToolOutput[OutputDataT] = str,
+        output_type: _output.OutputType[OutputDataT] = str,
         instructions: str
         | _system_prompt.SystemPromptFunc[AgentDepsT]
         | Sequence[str | _system_prompt.SystemPromptFunc[AgentDepsT]]
@@ -199,7 +199,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         name: str | None = None,
         model_settings: ModelSettings | None = None,
         retries: int = 1,
-        result_tool_name: str = 'final_result',
+        result_tool_name: str = _output.DEFAULT_OUTPUT_TOOL_NAME,
         result_tool_description: str | None = None,
         result_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
@@ -214,7 +214,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         model: models.Model | models.KnownModelName | str | None = None,
         *,
-        # TODO change this back to `output_type: type[OutputDataT] | ToolOutput[OutputDataT] = str,` when we remove the overloads
+        # TODO change this back to `output_type: _output.OutputType[OutputDataT] = str,` when we remove the overloads
         output_type: Any = str,
         instructions: str
         | _system_prompt.SystemPromptFunc[AgentDepsT]
@@ -374,7 +374,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT],
+        output_type: _output.OutputType[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -404,7 +404,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT] | None = None,
+        output_type: _output.OutputType[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -492,7 +492,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT],
+        output_type: _output.OutputType[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -524,7 +524,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT] | None = None,
+        output_type: _output.OutputType[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -770,7 +770,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT] | None,
+        output_type: _output.OutputType[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -800,7 +800,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT] | None = None,
+        output_type: _output.OutputType[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -883,7 +883,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent],
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT],
+        output_type: _output.OutputType[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -914,7 +914,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT] | None = None,
+        output_type: _output.OutputType[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -994,7 +994,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                                 if isinstance(maybe_part_event, _messages.PartStartEvent):
                                     new_part = maybe_part_event.part
                                     if isinstance(new_part, _messages.TextPart):
-                                        if _agent_graph.allow_text_output(output_schema):
+                                        if _output.allow_text_output(output_schema):
                                             return FinalResult(s, None, None)
                                     elif isinstance(new_part, _messages.ToolCallPart) and output_schema:
                                         for call, _ in output_schema.find_tool([new_part]):
@@ -1628,7 +1628,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         raise AttributeError('The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.')
 
     def _prepare_output_schema(
-        self, output_type: type[RunOutputDataT] | ToolOutput[RunOutputDataT] | None
+        self, output_type: _output.OutputType[RunOutputDataT] | None
     ) -> _output.OutputSchema[RunOutputDataT] | None:
         if output_type is not None:
             if self._output_validators:

pydantic_ai/mcp.py

@@ -373,7 +373,7 @@ class MCPServerHTTP(MCPServer):
             url=self.url,
             headers=self.headers,
             timeout=timedelta(seconds=self.timeout),
-            sse_read_timeout=timedelta(self.sse_read_timeout),
+            sse_read_timeout=timedelta(seconds=self.sse_read_timeout),
         ) as (read_stream, write_stream, _):
             yield read_stream, write_stream
 

pydantic_ai/messages.py

@@ -378,7 +378,7 @@ class ToolReturnPart:
         """Return a dictionary representation of the content, wrapping non-dict types appropriately."""
         # gemini supports JSON dict return values, but no other JSON types, hence we wrap anything else in a dict
         if isinstance(self.content, dict):
-            return tool_return_ta.dump_python(self.content, mode='json')  # pyright: ignore[reportUnknownMemberType]  # pragma: no cover
+            return tool_return_ta.dump_python(self.content, mode='json')  # pyright: ignore[reportUnknownMemberType]
         else:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
@@ -589,7 +589,7 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
-    vendor_details: dict[str, Any] | None = field(default=None, repr=False)
+    vendor_details: dict[str, Any] | None = field(default=None)
     """Additional vendor-specific details in a serializable format.
 
     This allows storing selected vendor-specific data that isn't mapped to standard ModelResponse fields.