refactor: address review feedback — processor validate/call, types, imports

Major changes:
- Add validate() and call() methods to BaseOutputProcessor with identity
  defaults. Override in ObjectOutputProcessor (existing), UnionOutputProcessor
  (new: clean decomposition of union envelope + inner validation/execution),
  TextFunctionOutputProcessor (new: function call in call()).
- Eliminate isinstance branching in _build_output_handlers — now uses
  polymorphic processor.validate() and processor.call().
- Move OutputContext to public pydantic_ai.output module.
- Fix output_type for functions: now the function input type (what the
  model produces), not the return type.
- Fix execute hook parameter types: Any instead of RawOutput, since
  validated output varies by processor type.
- Top-level imports in _tool_manager.py, keyword-only args in helpers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: DouweM

pydantic_ai_slim/pydantic_ai/_output.py

@@ -14,6 +14,7 @@
 
 from . import messages as _messages
 from ._instrumentation import InstrumentationNames
+from ._output import OutputToolset, run_output_execute_hooks, run_output_validate_hooks
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import (
     ApprovalRequired,
@@ -506,12 +507,6 @@ async def _raw_execute_output_tool(
         validated: ValidatedToolCall[AgentDepsT],
     ) -> Any:
         """Execute an output tool call with output hooks."""
-        from ._output import (
-            OutputToolset,
-            run_output_execute_hooks,
-            run_output_validate_hooks,
-        )
-
         assert validated.tool is not None
         assert validated.validated_args is not None
         assert self.root_capability is not None
@@ -535,18 +530,18 @@ async def do_validate(data: str | dict[str, Any]) -> str | dict[str, Any]:
 
         validated_output = await run_output_validate_hooks(
             cap,
-            ctx,
-            output_context,
-            validated.validated_args,
-            do_validate,
+            run_context=ctx,
+            output_context=output_context,
+            raw_output=validated.validated_args,
+            do_validate=do_validate,
             allow_partial=False,
             wrap_validation_errors=False,
         )
 
         # --- Output execute phase (wraps processor.call + output validators) ---
-        async def do_execute(output: str | dict[str, Any]) -> Any:
+        async def do_execute(output: Any) -> Any:
             try:
-                result = await processor.call(output, ctx, wrap_validation_errors=False)  # type: ignore[arg-type]
+                result = await processor.call(output, ctx, wrap_validation_errors=False)
                 for validator in toolset.output_validators:
                     result = await validator.validate(result, ctx, wrap_validation_errors=False)
                 return result
@@ -555,7 +550,13 @@ async def do_execute(output: str | dict[str, Any]) -> Any:
                 self.failed_tools.add(name)
                 raise self._wrap_error_as_retry(name, validated.call, e) from e
 
-        return await run_output_execute_hooks(cap, ctx, output_context, validated_output, do_execute)
+        return await run_output_execute_hooks(
+            cap,
+            run_context=ctx,
+            output_context=output_context,
+            validated=validated_output,
+            do_execute=do_execute,
+        )
 
     async def _execute_function_tool_call(
         self,

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from pydantic_ai._output import OutputContext
+from pydantic_ai.output import OutputContext
 
 from .abstract import (
     AbstractCapability,

pydantic_ai_slim/pydantic_ai/capabilities/__init__.py

@@ -15,10 +15,10 @@
 
 if TYPE_CHECKING:
     from pydantic_ai import _agent_graph
-    from pydantic_ai._output import OutputContext
     from pydantic_ai.agent.abstract import AgentModelSettings
     from pydantic_ai.capabilities.prefix_tools import PrefixTools
     from pydantic_ai.models import ModelRequestContext
+    from pydantic_ai.output import OutputContext
     from pydantic_ai.result import FinalResult
     from pydantic_ai.run import AgentRunResult
     from pydantic_graph import End
@@ -59,7 +59,7 @@
 WrapOutputValidateHandler: TypeAlias = 'Callable[[str | dict[str, Any]], Awaitable[str | dict[str, Any]]]'
 """Handler type for wrap_output_validate."""
 
-WrapOutputExecuteHandler: TypeAlias = 'Callable[[str | dict[str, Any]], Awaitable[Any]]'
+WrapOutputExecuteHandler: TypeAlias = 'Callable[[Any], Awaitable[Any]]'
 """Handler type for wrap_output_execute."""
 
 
@@ -570,17 +570,17 @@ async def before_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
-    ) -> RawOutput:
+    ) -> Any:
         """Modify validated output before execution (extraction + function call)."""
         return output
 
     async def after_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        validated_output: RawOutput,
+        validated_output: Any,
         output: Any,
         output_context: OutputContext,
     ) -> Any:
@@ -591,7 +591,7 @@ async def wrap_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         handler: WrapOutputExecuteHandler,
     ) -> Any:
@@ -602,7 +602,7 @@ async def on_output_execute_error(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         error: Exception,
     ) -> Any:

pydantic_ai_slim/pydantic_ai/capabilities/abstract.py

@@ -20,8 +20,8 @@
 
 if TYPE_CHECKING:
     from pydantic_ai import _agent_graph
-    from pydantic_ai._output import OutputContext
     from pydantic_ai.models import ModelRequestContext
+    from pydantic_ai.output import OutputContext
     from pydantic_ai.result import FinalResult
     from pydantic_ai.run import AgentRunResult
     from pydantic_graph import End
@@ -462,9 +462,9 @@ async def before_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
-    ) -> RawOutput:
+    ) -> Any:
         for capability in self.capabilities:
             output = await capability.before_output_execute(ctx, output=output, output_context=output_context)
         return output
@@ -473,7 +473,7 @@ async def after_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        validated_output: RawOutput,
+        validated_output: Any,
         output: Any,
         output_context: OutputContext,
     ) -> Any:
@@ -487,7 +487,7 @@ async def wrap_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         handler: WrapOutputExecuteHandler,
     ) -> Any:
@@ -500,7 +500,7 @@ async def on_output_execute_error(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         error: Exception,
     ) -> Any:
@@ -606,9 +606,9 @@ def _make_output_execute_wrap(
     cap: AbstractCapability[AgentDepsT],
     ctx: RunContext[AgentDepsT],
     output_context: OutputContext,
-    inner: Callable[[RawOutput], Awaitable[Any]],
-) -> Callable[[RawOutput], Awaitable[Any]]:
-    async def wrapped(output: RawOutput) -> Any:
+    inner: Callable[[Any], Awaitable[Any]],
+) -> Callable[[Any], Awaitable[Any]]:
+    async def wrapped(output: Any) -> Any:
         return await cap.wrap_output_execute(ctx, output=output, output_context=output_context, handler=inner)
 
     return wrapped

pydantic_ai_slim/pydantic_ai/capabilities/combined.py

@@ -48,8 +48,8 @@ async def log_request(ctx, request_context):
 )
 
 if TYPE_CHECKING:
-    from pydantic_ai._output import OutputContext
     from pydantic_ai.models import ModelRequestContext
+    from pydantic_ai.output import OutputContext
     from pydantic_ai.run import AgentRunResult
 
 _FuncT = TypeVar('_FuncT', bound=Callable[..., Any])
@@ -202,19 +202,19 @@ def __call__(self, ctx: RunContext[Any], /, *, raw_output: RawOutput, output_con
 
 class BeforeOutputExecuteHookFunc(Protocol):
     """Protocol for :meth:`~AbstractCapability.before_output_execute` hook functions."""
-    def __call__(self, ctx: RunContext[Any], /, *, output: RawOutput, output_context: OutputContext) -> RawOutput | Awaitable[RawOutput]: ...
+    def __call__(self, ctx: RunContext[Any], /, *, output: Any, output_context: OutputContext) -> Any | Awaitable[Any]: ...
 
 class AfterOutputExecuteHookFunc(Protocol):
     """Protocol for :meth:`~AbstractCapability.after_output_execute` hook functions."""
-    def __call__(self, ctx: RunContext[Any], /, *, validated_output: RawOutput, output: Any, output_context: OutputContext) -> Any | Awaitable[Any]: ...
+    def __call__(self, ctx: RunContext[Any], /, *, validated_output: Any, output: Any, output_context: OutputContext) -> Any | Awaitable[Any]: ...
 
 class WrapOutputExecuteHookFunc(Protocol):
     """Protocol for :meth:`~AbstractCapability.wrap_output_execute` hook functions."""
-    def __call__(self, ctx: RunContext[Any], /, *, output: RawOutput, output_context: OutputContext, handler: WrapOutputExecuteHandler) -> Any | Awaitable[Any]: ...
+    def __call__(self, ctx: RunContext[Any], /, *, output: Any, output_context: OutputContext, handler: WrapOutputExecuteHandler) -> Any | Awaitable[Any]: ...
 
 class OnOutputExecuteErrorHookFunc(Protocol):
     """Protocol for :meth:`~AbstractCapability.on_output_execute_error` hook functions."""
-    def __call__(self, ctx: RunContext[Any], /, *, output: RawOutput, output_context: OutputContext, error: Exception) -> Any | Awaitable[Any]: ...
+    def __call__(self, ctx: RunContext[Any], /, *, output: Any, output_context: OutputContext, error: Exception) -> Any | Awaitable[Any]: ...
 # fmt: on
 
 
@@ -1114,16 +1114,16 @@ async def on_output_validate_error(
         raise error
 
     async def before_output_execute(
-        self, ctx: RunContext[AgentDepsT], *, output: RawOutput, output_context: OutputContext
-    ) -> RawOutput:
+        self, ctx: RunContext[AgentDepsT], *, output: Any, output_context: OutputContext
+    ) -> Any:
         for entry in self._get('before_output_execute'):
             output = await _call_entry(
                 entry, 'before_output_execute', ctx, output=output, output_context=output_context
             )
         return output
 
     async def after_output_execute(
-        self, ctx: RunContext[AgentDepsT], *, validated_output: RawOutput, output: Any, output_context: OutputContext
+        self, ctx: RunContext[AgentDepsT], *, validated_output: Any, output: Any, output_context: OutputContext
     ) -> Any:
         for entry in self._get('after_output_execute'):
             output = await _call_entry(
@@ -1140,7 +1140,7 @@ async def wrap_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         handler: WrapOutputExecuteHandler,
     ) -> Any:
@@ -1158,7 +1158,7 @@ async def on_output_execute_error(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         error: Exception,
     ) -> Any:

pydantic_ai_slim/pydantic_ai/capabilities/hooks.py

@@ -29,9 +29,9 @@
 )
 
 if TYPE_CHECKING:
-    from pydantic_ai._output import OutputContext
     from pydantic_ai.agent.abstract import AgentModelSettings
     from pydantic_ai.models import ModelRequestContext
+    from pydantic_ai.output import OutputContext
     from pydantic_ai.run import AgentRunResult
 
 
@@ -338,16 +338,16 @@ async def before_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
-    ) -> RawOutput:
+    ) -> Any:
         return await self.wrapped.before_output_execute(ctx, output=output, output_context=output_context)
 
     async def after_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        validated_output: RawOutput,
+        validated_output: Any,
         output: Any,
         output_context: OutputContext,
     ) -> Any:
@@ -359,7 +359,7 @@ async def wrap_output_execute(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         handler: WrapOutputExecuteHandler,
     ) -> Any:
@@ -371,7 +371,7 @@ async def on_output_execute_error(
         self,
         ctx: RunContext[AgentDepsT],
         *,
-        output: RawOutput,
+        output: Any,
         output_context: OutputContext,
         error: Exception,
     ) -> Any:

pydantic_ai_slim/pydantic_ai/capabilities/wrapper.py

@@ -23,6 +23,7 @@
     'TextOutput',
     'StructuredDict',
     'OutputObjectDefinition',
+    'OutputContext',
     # types
     'OutputDataT',
     'OutputMode',
@@ -268,6 +269,24 @@ class OutputObjectDefinition:
     strict: bool | None = None
 
 
+@dataclass
+class OutputContext:
+    """Context about the output being processed, passed to output hooks."""
+
+    mode: OutputMode
+    """The output mode ('text', 'native', 'prompted', 'tool', 'auto')."""
+    output_type: type[Any] | None
+    """The resolved output type (e.g. MyModel, str). For output functions, the function's input type (what the model produces)."""
+    object_def: OutputObjectDefinition | None
+    """The output object definition (schema, name, description), if structured output."""
+    has_function: bool
+    """Whether there's an output function to call in the execute step."""
+    tool_call: ToolCallPart | None = None
+    """The tool call part, for tool-based output. None for text output."""
+    tool_def: ToolDefinition | None = None
+    """The tool definition, for tool-based output. None for text output."""
+
+
 @dataclass
 class TextOutput(Generic[OutputDataT]):
     """Marker class to use text output for an output function taking a string argument.

pydantic_ai_slim/pydantic_ai/output.py

@@ -8,7 +8,6 @@
 import pytest
 from pydantic import BaseModel, ValidationError
 
-from pydantic_ai._output import OutputContext
 from pydantic_ai._run_context import RunContext
 from pydantic_ai.agent import Agent
 from pydantic_ai.capabilities.abstract import AbstractCapability
@@ -21,7 +20,7 @@
     ToolCallPart,
 )
 from pydantic_ai.models.function import AgentInfo, FunctionModel
-from pydantic_ai.output import PromptedOutput, TextOutput
+from pydantic_ai.output import OutputContext, PromptedOutput, TextOutput
 from pydantic_ai.tools import ToolDefinition
 
 pytestmark = [