Allow typed `RunContext[Deps]` in `TextOutput` signature #3732

dsfaccini · 2025-12-15T16:07:35Z

TextOutput supports passing in functions that take in the agent's RunContext[Deps], but pyright complains when the Deps type is specified in the signature. This PR adds support for properly typing the RunContext

Context (badum tss): https://pydantic.slack.com/archives/C081LUCJ4KX/p1765814128383649?thread_ts=1765812544.296369&cid=C081LUCJ4KX

Note: sorry about the change in the uv.lock, that seems to have escaped the google-genai bump PR

samuelcolvin

LGTM

github-actions · 2025-12-15T16:22:57Z

Docs Preview

commit:	`2d3eb6f`
Preview URL:	https://dc7aeeac-pydantic-ai-previews.pydantic.workers.dev

docs/output.md

DouweM · 2025-12-15T19:09:40Z

pydantic_ai_slim/pydantic_ai/output.py

    '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],


See #3319 and #3721; we shouldn't be reusing generic type vars across different files, but instead need a new one here.
And we should think about covariant vs contravariant etc.

The easiest way to ensure it works as expected is to add some cases that should be valid to typed_agent.py, which is automatically type checked

DouweM · 2025-12-15T19:10:31Z

tests/test_agent.py

    )


+def test_output_type_text_output_function_with_deps():


I'm pretty sure we've already tested somewhere that it works, this PR is just for fixing the type checking, so let's remove the new tests and add some cases to typed_agent.py instead.

I think we can remove this test

uv.lock

…ture

DouweM · 2025-12-16T14:21:45Z

docs/output.md

-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.

-```python {title="text_output_function.py"}


Sorry I didn't mean to remove the example entirely; just the extra RunContext-based on you'd added 😄

lmao I was like "okay..."

DouweM · 2025-12-16T14:23:17Z

tests/test_examples.py

    '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',


Not needed anymore

DouweM · 2025-12-16T14:25:42Z

tests/typed_agent.py

 Agent('test', tools=[Tool(foobar_plain)], deps_type=int)  # pyright: ignore[reportArgumentType,reportCallIssue]

+# TextOutput with RunContext
+Agent('test', output_type=TextOutput(str_to_regex_with_ctx), deps_type=int)


To verify that contravariant=True is the right option, we should also have an example where the deps_type is a subclass of the type on the function's ctx, e.g. bool. That should work, because a function that takes an int is also able to handle a bool (which is essentially 1 or 0.

Can we move TextOutput(str_to_regex_with_ctx) to a variable and add an assert_type so we can verify its type is inferred correctly?

DouweM · 2025-12-16T14:26:59Z

tests/typed_agent.py

+# TextOutput with RunContext
+Agent('test', output_type=TextOutput(str_to_regex_with_ctx), deps_type=int)
+Agent('test', output_type=TextOutput(str_to_regex_with_ctx), deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
+Agent('test', output_type=TextOutput(str_to_regex_with_ctx))  # pyright: ignore[reportArgumentType,reportCallIssue]


We can then reuse the var here

In _OutputSpecItem, we fill in the generic param as Any, so I'm curious if that could cause any issues here because maybe it'll accept any type instead of just same one as deps_type? It'd be interesting to see the typing errors you got here before you add the pyright: ignores. (just pasting them into a GH comment here is fine)

no type errors, I added a NOTE for it

…ture

DouweM · 2025-12-17T17:46:52Z

docs/output.md


 _(This example is complete, it can be run "as is")_

+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).


Above the example please :) I prefer to explain, then show

DouweM · 2025-12-17T17:47:38Z

tests/test_agent.py

    )


+def test_output_type_text_output_function_with_deps():


I think we can remove this test

DouweM · 2025-12-17T17:52:47Z

tests/typed_agent.py

+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.


If we can't make those give a typing error, then maybe it's not worth adding TextOutputAgentDepsT and we can just make it RunContext[Any].

Ideally, we would go a step further and add DepsT to _OutputSpecItem, OutputSpec etc, so that the user is warned if their ctx: RunContext[Foo] arg doesn't match the agent's deps_type, which would work for TextOutput(func), but unfortunately there's now way to make that work for output functions in general, e.g. output_type=func or output_type=ToolOutput(func), as those functions that arbitrary args, so if the first arg ctx: RunContext[Foo] while deps_type=Bar, the type checker will just treat it like an arbitrary arg, instead of complaining. (ping me if that doesn't make sense; typing is tricky)

So I think the best we can do is just make it RunContext[Any] and accept that this part isn't type checked.

I think I was able to enforce this

Co-authored-by: Douwe Maan <[email protected]>

DouweM · 2025-12-17T21:37:01Z

pydantic_ai_slim/pydantic_ai/output.py

    '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],


Yeah looking at this now, since this will only work for TextOutput and not output functions in general, I don't think it's worth the extra (out of order) type var and would prefer to just have RunContext[Any] here and leave it at that

DouweM · 2025-12-18T22:04:20Z

pydantic_ai_slim/pydantic_ai/output.py


 @dataclass
-class TextOutput(Generic[OutputDataT]):
+class TextOutput(Generic[OutputDataT, TextOutputAgentDepsT]):


Let's drop TextOutputAgentDepsT entirely and use RunContext[Any] in TextOutputFunc, as we don't actually use the generic param and it bothers me that it's out of order compared to everything else that takes Deps followed by Output 😄

DouweM · 2025-12-18T22:04:57Z

tests/typed_agent.py

+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]


I agree it's worth having a comment like this, but this'll have to be rewritten if my comment above is addressed

dsfaccini added 2 commits December 15, 2025 11:04

support

e84ab36

fix example test

cc5023f

samuelcolvin approved these changes Dec 15, 2025

View reviewed changes

DouweM requested changes Dec 15, 2025

View reviewed changes

DouweM self-assigned this Dec 15, 2025

DouweM added the awaiting author revision label Dec 15, 2025

dsfaccini added 5 commits December 15, 2025 17:15

support

248faa6

fix example test

9e0ab5b

requedsted changed - pending

c00d85c

requested changes - pending

becf787

Merge branch 'main' into support-typed-runcontext-in-textoutput-signa…

fab8356

…ture

DouweM requested changes Dec 16, 2025

View reviewed changes

dsfaccini added 2 commits December 17, 2025 12:42

requested changes

b6ae5b9

Merge branch 'main' into support-typed-runcontext-in-textoutput-signa…

5d15e4b

…ture

DouweM requested changes Dec 17, 2025

View reviewed changes

dsfaccini and others added 2 commits December 17, 2025 13:03

Apply suggestions from code review

189dce5

Co-authored-by: Douwe Maan <[email protected]>

remove test and duplicate text

9712a00

DouweM requested changes Dec 17, 2025

View reviewed changes

dsfaccini force-pushed the support-typed-runcontext-in-textoutput-signature branch from 345bc95 to 9712a00 Compare December 18, 2025 16:12

link to PR comment

3d3e337

DouweM requested changes Dec 18, 2025

View reviewed changes

removed TextOutputAgentDepsT

2d3eb6f


		_(This example is complete, it can be run "as is")_

		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).

		)


		def test_output_type_text_output_function_with_deps():

		)


		def test_output_type_text_output_function_with_deps():

Allow typed RunContext[Deps] in TextOutput signature #3732

Are you sure you want to change the base?

Allow typed RunContext[Deps] in TextOutput signature #3732

Conversation

dsfaccini commented Dec 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

samuelcolvin left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Dec 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Docs Preview

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Allow typed `RunContext[Deps]` in `TextOutput` signature #3732

Allow typed `RunContext[Deps]` in `TextOutput` signature #3732

dsfaccini commented Dec 15, 2025 •

edited

Loading

github-actions bot commented Dec 15, 2025 •

edited

Loading