Allow typed RunContext[Deps] in TextOutput signature

docs/output.md

@@ -237,7 +237,11 @@ RouterFailure(explanation='I am not equipped to provide travel information, such
 
 #### Text output
 
-If you provide an output function that takes a string, Pydantic AI will by default create an output tool like for any other output function. If instead you'd like the model to provide the string using plain text output, you can wrap the function in the [`TextOutput`][pydantic_ai.output.TextOutput] marker class. If desired, this marker class can be used alongside one or more [`ToolOutput`](#tool-output) marker classes (or unmarked types or functions) in a list provided to `output_type`.
+If you provide an output function that takes a string, Pydantic AI will by default create an output tool like for any other output function. If instead you'd like the model to provide the string using plain text output, you can wrap the function in the [`TextOutput`][pydantic_ai.output.TextOutput] marker class.
+
+If desired, this marker class can be used alongside one or more [`ToolOutput`](#tool-output) marker classes (or unmarked types or functions) in a list provided to `output_type`.
+
+Like other output functions, text output functions can optionally take [`RunContext`][pydantic_ai.tools.RunContext] as the first argument, and can raise [`ModelRetry`][pydantic_ai.exceptions.ModelRetry] to ask the model to try again with modified arguments (or with a different output type).
 
 ```python {title="text_output_function.py"}
 from pydantic_ai import Agent, TextOutput

pydantic_ai_slim/pydantic_ai/output.py

@@ -11,8 +11,9 @@
 
 from . import _utils, exceptions
 from ._json_schema import InlineDefsJsonSchemaTransformer
+from ._run_context import RunContext
 from .messages import ToolCallPart
-from .tools import DeferredToolRequests, ObjectJsonSchema, RunContext, ToolDefinition
+from .tools import DeferredToolRequests, ObjectJsonSchema, ToolDefinition
 
 __all__ = (
     # classes
@@ -33,6 +34,7 @@
 
 T = TypeVar('T')
 T_co = TypeVar('T_co', covariant=True)
+TextOutputAgentDepsT = TypeVar('TextOutputAgentDepsT', default=None, contravariant=True)
 
 OutputDataT = TypeVar('OutputDataT', default=str, covariant=True)
 """Covariant type variable for the output data type of a run."""
@@ -60,8 +62,8 @@
 
 TextOutputFunc = TypeAliasType(
     'TextOutputFunc',
-    Callable[[RunContext, str], Awaitable[T_co] | T_co] | Callable[[str], Awaitable[T_co] | T_co],
-    type_params=(T_co,),
+    Callable[[RunContext[TextOutputAgentDepsT], str], Awaitable[T_co] | T_co] | Callable[[str], Awaitable[T_co] | T_co],
+    type_params=(T_co, TextOutputAgentDepsT),
 )
 """Definition of a function that will be called to process the model's plain text output. The function must take a single string argument.
 
@@ -259,7 +261,7 @@ class OutputObjectDefinition:
 
 
 @dataclass
-class TextOutput(Generic[OutputDataT]):
+class TextOutput(Generic[OutputDataT, TextOutputAgentDepsT]):
     """Marker class to use text output for an output function taking a string argument.
 
     Example:
@@ -281,7 +283,7 @@ def split_into_words(text: str) -> list[str]:
     ```
     """
 
-    output_function: TextOutputFunc[OutputDataT]
+    output_function: TextOutputFunc[OutputDataT, TextOutputAgentDepsT]
     """The function that will be called to process the model's plain text output. The function must take a single string argument."""
 
 
@@ -354,7 +356,7 @@ def __get_pydantic_json_schema__(
 
 _OutputSpecItem = TypeAliasType(
     '_OutputSpecItem',
-    OutputTypeOrFunction[T_co] | ToolOutput[T_co] | NativeOutput[T_co] | PromptedOutput[T_co] | TextOutput[T_co],
+    OutputTypeOrFunction[T_co] | ToolOutput[T_co] | NativeOutput[T_co] | PromptedOutput[T_co] | TextOutput[T_co, Any],
     type_params=(T_co,),
 )
 

tests/typed_agent.py

@@ -196,6 +196,10 @@ def str_to_regex(text: str) -> re.Pattern[str]:
     return re.compile(text)
 
 
+def str_to_regex_with_ctx(ctx: RunContext[int], text: str) -> re.Pattern[str]:
+    return re.compile(text)
+
+
 class MyClass:
     def my_method(self) -> bool:
         return True
@@ -283,6 +287,17 @@ def my_method(self) -> bool:
 # since deps are not set, they default to `None`, so can't be `int`
 Agent('test', tools=[Tool(foobar_plain)], deps_type=int)  # pyright: ignore[reportArgumentType,reportCallIssue]
 
+# TextOutput with RunContext
+text_output_with_ctx = TextOutput(str_to_regex_with_ctx)
+assert_type(text_output_with_ctx, TextOutput[re.Pattern[str], int])
+Agent('test', output_type=text_output_with_ctx, deps_type=int)
+Agent('test', output_type=text_output_with_ctx, deps_type=bool)  # bool is subclass of int, works with contravariant
+# NOTE: The following don't produce type errors because _OutputSpecItem uses TextOutput[T_co, Any]
+# which erases the deps type constraint. This is a known limitation.
+# https://github.com/pydantic/pydantic-ai/pull/3732#discussion_r2628741424
+Agent('test', output_type=text_output_with_ctx, deps_type=str)
+Agent('test', output_type=text_output_with_ctx)
+
 # prepare example from docs: