Add partial_output to RunContext

[petersli]

What

Resolves #3194

[DouweM]

Checking the length like this means we only support partial when you're also taking ctx, which is not clear from the docstring or the docs. Can we make it work with def validator(output, partial) as well? We can use the _takes_ctx helper from _function_schema.py to actually check the type of the first arg.

[petersli]

yeah, can definitely support that. Doesn't seem like the _takes_ctx helper works here though, it calling get_function_type_hints results in NameError in test_output_tool_retry_error_handled

    def test_output_tool_retry_error_handled():
        class OutputModel(BaseModel):
            x: int
            y: str

        agent = Agent('test', output_type=OutputModel, retries=2)

        call_count = 0

>       @agent.output_validator

tests/models/test_model_test.py:259:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
pydantic_ai_slim/pydantic_ai/agent/__init__.py:1018: in output_validator
    self._output_validators.append(_output.OutputValidator[AgentDepsT, Any](func))
../../.local/share/uv/python/cpython-3.12.12-macos-aarch64-none/lib/python3.12/typing.py:1184: in __call__
    result = self.__origin__(*args, **kwargs)
<string>:4: in __init__
    ???
pydantic_ai_slim/pydantic_ai/_output.py:174: in __post_init__
    self._takes_ctx = _function_schema._takes_ctx(self.function)
pydantic_ai_slim/pydantic_ai/_function_schema.py:263: in _takes_ctx
    type_hints = _typing_extra.get_function_type_hints(_decorators.unwrap_wrapped_function(callable_obj))
.venv/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py:551: in get_function_type_hints
    type_hints[name] = eval_type_backport(value, globalns, localns, type_params)
.venv/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py:429: in eval_type_backport
    return _eval_type_backport(value, globalns, localns, type_params)
.venv/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py:466: in _eval_type_backport
    return _eval_type(value, globalns, localns, type_params)
.venv/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py:500: in _eval_type
    return typing._eval_type(  # type: ignore
../../.local/share/uv/python/cpython-3.12.12-macos-aarch64-none/lib/python3.12/typing.py:415: in _eval_type
    return t._evaluate(globalns, localns, type_params, recursive_guard=recursive_guard)
../../.local/share/uv/python/cpython-3.12.12-macos-aarch64-none/lib/python3.12/typing.py:947: in _evaluate
    eval(self.__forward_code__, globalns, localns),
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

>   ???
E   NameError: name 'OutputModel' is not defined

[DouweM]

Hmm, that's because OutputModel is defined in the local namespace rather than the global one. get_function_type_hints takes a localns arg but we're not passing it, and it'd be tough to pass it all the way through the stack to this call :/

I wonder if we could use this util: https://github.com/pydantic/pydantic/blob/de79df29547cede38415b0ae6c4d45fd0a9cab4e/pydantic/_internal/_namespace_utils.py#L106

[petersli]

hmm, the localns returned by that util still doesn't include OutputModel

I feel like the pragmatic thing to do is to avoid all this complexity and only support partial on validators that take ctx (as originally proposed), and just make that clear in the docs -- what do you think?

[petersli]

I reverted to the original supported signatures and updated the docs and docstring to make it clear

[DouweM]

This means that when streaming text, the validator will never be called without partial=True (unless the user explicitly reads the output later), meaning that validation may end up being skipped entirely. I think we should do what we do in stream_output() and do one final call after streaming with partial=False

[petersli]

hmm I implemented this, but then I realized it's actually not needed currently -- see test_output_validator_partial_stream_text, which verifies that the validator is called with partial=False. It happens as part of the on_complete in run_stream. If I add another validation call, it actually gets called twice.