Allow typed RunContext[Deps] in TextOutput signature

docs/output.md

@@ -240,9 +240,12 @@ RouterFailure(explanation='I am not equipped to provide travel information, such
 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`.
 
 ```python {title="text_output_function.py"}
-from pydantic_ai import Agent, TextOutput
+from dataclasses import dataclass
+
+from pydantic_ai import Agent, RunContext, TextOutput
 
 
+# Without deps
 def split_into_words(text: str) -> list[str]:
     return text.split()
 
@@ -254,6 +257,26 @@ agent = Agent(
 result = agent.run_sync('Who was Albert Einstein?')
 print(result.output)
 #> ['Albert', 'Einstein', 'was', 'a', 'German-born', 'theoretical', 'physicist.']
+
+
+# Or with deps
+@dataclass
+class Deps:
+    prefix: str
+
+
+def add_prefix_and_split(ctx: RunContext[Deps], text: str) -> list[str]:
+    return f'{ctx.deps.prefix}: {text}'.split()
+
+
+agent_with_deps = Agent(
+    'openai:gpt-5',
+    deps_type=Deps,
+    output_type=TextOutput(add_prefix_and_split),
+)
+result = agent_with_deps.run_sync('Hello world', deps=Deps(prefix='Response'))
+print(result.output)
+#> ['Response:', 'Hello', 'world']
 ```
 
 _(This example is complete, it can be run "as is")_

pydantic_ai_slim/pydantic_ai/output.py

@@ -11,8 +11,9 @@
 
 from . import _utils, exceptions
 from ._json_schema import InlineDefsJsonSchemaTransformer
+from ._run_context import AgentDepsT, RunContext
 from .messages import ToolCallPart
-from .tools import DeferredToolRequests, ObjectJsonSchema, RunContext, ToolDefinition
+from .tools import DeferredToolRequests, ObjectJsonSchema, ToolDefinition
 
 __all__ = (
     # classes
@@ -60,8 +61,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[AgentDepsT], str], Awaitable[T_co] | T_co] | Callable[[str], Awaitable[T_co] | T_co],
+    type_params=(T_co, AgentDepsT),
 )
 """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 +260,7 @@ class OutputObjectDefinition:
 
 
 @dataclass
-class TextOutput(Generic[OutputDataT]):
+class TextOutput(Generic[OutputDataT, AgentDepsT]):
     """Marker class to use text output for an output function taking a string argument.
 
     Example:
@@ -281,7 +282,7 @@ def split_into_words(text: str) -> list[str]:
     ```
     """
 
-    output_function: TextOutputFunc[OutputDataT]
+    output_function: TextOutputFunc[OutputDataT, AgentDepsT]
     """The function that will be called to process the model's plain text output. The function must take a single string argument."""
 
 
@@ -354,7 +355,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/test_agent.py

@@ -1167,6 +1167,24 @@ def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
     )
 
 
+def test_output_type_text_output_function_with_deps():
+    """Test that TextOutput functions can use RunContext with custom deps type."""
+
+    @dataclass
+    class Deps:
+        prefix: str
+
+    def add_prefix_and_split(ctx: RunContext[Deps], text: str) -> list[str]:
+        return f'{ctx.deps.prefix}: {text}'.split()
+
+    def return_text(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        return ModelResponse(parts=[TextPart(content='Hello world')])
+
+    agent = Agent(FunctionModel(return_text), deps_type=Deps, output_type=TextOutput(add_prefix_and_split))
+    result = agent.run_sync('test prompt', deps=Deps(prefix='Response'))
+    assert result.output == snapshot(['Response:', 'Hello', 'world'])
+
+
 @pytest.mark.parametrize(
     'output_type',
     [[str, str], [str, TextOutput(upcase)], [TextOutput(upcase), TextOutput(upcase)]],

tests/test_examples.py

@@ -344,6 +344,7 @@ async def call_tool(
     'What is the capital of the UK?': 'The capital of the UK is London.',
     'What is the capital of Mexico?': 'The capital of Mexico is Mexico City.',
     'Who was Albert Einstein?': 'Albert Einstein was a German-born theoretical physicist.',
+    'Hello world': 'Hello world',
     'What was his most famous equation?': "Albert Einstein's most famous equation is (E = mc^2).",
     'What is the date?': 'Hello Frank, the date today is 2032-01-02.',
     'What is this? https://ai.pydantic.dev': 'A Python agent framework for building Generative AI applications.',