simplify: only support (ctx, output, partial) signature

Author: petersli

docs/output.md

@@ -474,15 +474,15 @@ _(This example is complete, it can be run "as is")_
 
 When [streaming responses](#streaming-model-responses), output validators are called multiple times - once for each partial response and once for the final response. By default, validators receive `allow_partial=False`, meaning they treat all responses the same way.
 
-However, you can add a `partial: bool` parameter to your validator to distinguish between partial and final validation. This is useful when you want to skip expensive validation during streaming but apply full validation to the final result:
+However, you can add a `partial: bool` parameter as the last argument to your validator to distinguish between partial and final validation. This is useful when you want to skip expensive validation during streaming but apply full validation to the final result:
 
 ```python
 from pydantic_ai import Agent, ModelRetry
 
 agent = Agent('openai:gpt-4o')
 
 @agent.output_validator
-def validate_output(output: str, *, partial: bool) -> str:
+def validate_output(ctx: RunContext, output: str, partial: bool) -> str:
     if partial:
         return output
     else:
@@ -491,13 +491,6 @@ def validate_output(output: str, *, partial: bool) -> str:
         return output
 ```
 
-The `partial` parameter must be keyword-only (note the `*` before `partial`). It works with all validator signatures:
-
-- `(data: T) -> T` - simple validator, no partial awareness
-- `(data: T, *, partial: bool) -> T` - validator with partial parameter
-- `(ctx: RunContext[Deps], data: T) -> T` - validator with context
-- `(ctx: RunContext[Deps], data: T, *, partial: bool) -> T` - validator with context and partial
-
 ## Image output
 
 Some models can generate images as part of their response, for example those that support the [Image Generation built-in tool](builtin-tools.md#image-generation-tool) and OpenAI models using the [Code Execution built-in tool](builtin-tools.md#code-execution-tool) when told to generate a chart.

pydantic_ai_slim/pydantic_ai/_output.py

@@ -53,23 +53,20 @@
 resolve these potential variance issues.
 """
 
-OutputValidatorFunc: TypeAlias = Callable[..., Any] | Callable[..., Awaitable[Any]]
+OutputValidatorFunc = (
+    Callable[[RunContext[AgentDepsT], OutputDataT_inv, bool], OutputDataT_inv]
+    | Callable[[RunContext[AgentDepsT], OutputDataT_inv, bool], Awaitable[OutputDataT_inv]]
+    | 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 takes and returns the same type of data (which is the result type of an agent run), and:
+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 take a keyword-only `partial: bool` parameter (e.g., `def validator(data: T, *, partial: bool)`)
+* may or may not take a `partial: bool` parameter as a last argument
 * may or may not be async
-
-Usage `OutputValidatorFunc[AgentDepsT, T]`.
-
-The function signature is introspected at runtime to determine which parameters it accepts.
-Supported signatures:
-- `(data: T) -> T`
-- `(data: T, *, partial: bool) -> T`
-- `(ctx: RunContext[Deps], data: T) -> T`
-- `(ctx: RunContext[Deps], data: T, *, partial: bool) -> T`
-- Async variants of all above
 """
 
 
@@ -171,25 +168,8 @@ class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
 
     def __post_init__(self):
         sig = inspect.signature(self.function)
-
-        if 'partial' in sig.parameters:
-            partial_param = sig.parameters['partial']
-            if partial_param.kind != inspect.Parameter.KEYWORD_ONLY:
-                raise ValueError(
-                    f'Output validator {self.function.__name__!r} has a `partial` parameter that is not keyword-only. '
-                    'The `partial` parameter must be keyword-only (e.g., `def validator(output: str, *, partial: bool)`).'
-                )
-            self._takes_partial = True
-        else:
-            self._takes_partial = False
-
-        positional_params = [
-            p
-            for p in sig.parameters.values()
-            if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
-        ]
-        self._takes_ctx = len(positional_params) > 1
-
+        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(
@@ -216,17 +196,15 @@ async def validate(
             args = (result,)
 
         if self._takes_partial:
-            kwargs = {'partial': allow_partial}
-        else:
-            kwargs = {}
+            args = (*args, allow_partial)
 
         try:
             if self._is_async:
                 function = cast(Callable[[Any], Awaitable[T]], self.function)
-                result_data = await function(*args, **kwargs)
+                result_data = await function(*args)
             else:
                 function = cast(Callable[[Any], T], self.function)
-                result_data = await _utils.run_in_executor(function, *args, **kwargs)
+                result_data = await _utils.run_in_executor(function, *args)
         except ModelRetry as r:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -951,7 +951,16 @@ def decorator(
             self._system_prompt_functions.append(_system_prompt.SystemPromptRunner[AgentDepsT](func, dynamic=dynamic))
             return func
 
-    # Without partial parameter
+    @overload
+    def output_validator(
+        self, func: Callable[[RunContext[AgentDepsT], OutputDataT, bool], OutputDataT], /
+    ) -> Callable[[RunContext[AgentDepsT], OutputDataT, bool], OutputDataT]: ...
+
+    @overload
+    def output_validator(
+        self, func: Callable[[RunContext[AgentDepsT], OutputDataT, bool], Awaitable[OutputDataT]], /
+    ) -> Callable[[RunContext[AgentDepsT], OutputDataT, bool], Awaitable[OutputDataT]]: ...
+
     @overload
     def output_validator(
         self, func: Callable[[RunContext[AgentDepsT], OutputDataT], OutputDataT], /
@@ -972,13 +981,9 @@ def output_validator(
         self, func: Callable[[OutputDataT], Awaitable[OutputDataT]], /
     ) -> Callable[[OutputDataT], Awaitable[OutputDataT]]: ...
 
-    # With partial parameter (these use Any to bypass Callable's limitation with keyword-only params)
-    @overload
-    def output_validator(self, func: Any, /) -> Any: ...
-
     def output_validator(
-        self, func: _output.OutputValidatorFunc, /
-    ) -> _output.OutputValidatorFunc:
+        self, func: _output.OutputValidatorFunc[AgentDepsT, OutputDataT], /
+    ) -> _output.OutputValidatorFunc[AgentDepsT, OutputDataT]:
         """Decorator to register an output validator function.
 
         Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument.