Repository: pydantic/pydantic-ai
File: pydantic_ai_slim/pydantic_ai/_output.py (Python)

Code before:
            else:
                raise

        # Record response if content inclusion is enabled
        if run_context.trace_include_content and span.is_recording():
            from .models.instrumented import InstrumentedModel

            span.set_attribute(
                instrumentation_names.tool_result_attr,
                output if isinstance(output, str) else json.dumps(InstrumentedModel.serialize_any(output)),
            )

        return output


@dataclass
class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
    function: OutputValidatorFunc[AgentDepsT, OutputDataT_inv]
    _takes_ctx: bool = field(init=False)
    _takes_partial: bool = field(init=False)
    _is_async: bool = field(init=False)

    def __post_init__(self):
        sig = inspect.signature(self.function)
        self._takes_ctx = len(sig.parameters) > 1
        self._takes_partial = len(sig.parameters) > 2
        self._is_async = _utils.is_async_callable(self.function)

    async def validate(
        self,
        result: T,
        run_context: RunContext[AgentDepsT],
        wrap_validation_errors: bool = True,
        allow_partial: bool = False,
    ) -> T:
        """Validate a result but calling the function.

        Args:
            result: The result data after Pydantic validation the message content.
            run_context: The current run context.
            wrap_validation_errors: If true, wrap the validation errors in a retry message.
            allow_partial: Whether partial validation is allowed (passed to validator if it accepts it).

        Returns:
            Result of either the validated result data (ok) or a retry message (Err).
        """
        if self._takes_ctx:
            args = run_context, result
        else:
            args = (result,)


Review comment (refactor):
Checking the length like this means we only support `partial` when you're also taking `ctx`, which is not clear from the docstring or the docs. Can we make it work with `def validator(output, partial)` as well? We can use the `_takes_ctx` helper from `_function_schema.py` to actually check the type of the first arg.

Code after:
        # Record response if content inclusion is enabled
        if run_context.trace_include_content and span.is_recording():
            from .models.instrumented import InstrumentedModel

            span.set_attribute(
                instrumentation_names.tool_result_attr,
                output if isinstance(output, str) else json.dumps(InstrumentedModel.serialize_any(output)),
            )

        return output


@dataclass
class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
    function: OutputValidatorFunc[AgentDepsT, OutputDataT_inv]
    _takes_ctx: bool = field(init=False)
    _is_async: bool = field(init=False)

    def __post_init__(self):
        self._takes_ctx = len(inspect.signature(self.function).parameters) > 1
        self._is_async = _utils.is_async_callable(self.function)

    async def validate(
        self,
        result: T,
        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.
            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:
            args = run_context, result
        else:
            args = (result,)

        try:
            if self._is_async:
                function = cast(Callable[[Any], Awaitable[T]], self.function)
                result_data = await function(*args)
            else:
                function = cast(Callable[[Any], T], self.function)
                result_data = await _utils.run_in_executor(function, *args)
        except ModelRetry as r: