Raising ModelRetry in output validator function crashes when streaming

[petersli]

Initial Checks

• I confirm that I'm using the latest version of Pydantic AI
• I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue

Description

When using stream_output, if I raise ModelRetry in my output validator function, instead of seeing another model request (like I do when I raise ModelRetry in a function tool), the exception is unhandled.

Based on the docs and docstrings, I would not expect this to crash.

It works properly with run_sync, which is the case that's tested in tests/test_agent.py and the docs.

Example Code

# add this test to tests/test_agent.py
async def test_output_validator_stream_output():
    """Test that ModelRetry in output validators works correctly with streaming."""

    async def stream_model(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[DeltaToolCalls]:
        assert info.output_tools is not None
        if len(messages) == 1:
            # First attempt: return wrong value (a=41)
            yield {0: DeltaToolCall(name=info.output_tools[0].name, json_args='{"a": 41')}
            yield {0: DeltaToolCall(json_args=', "b": "f')}
            yield {0: DeltaToolCall(json_args='oo"}')}
        else:
            # Retry: return correct value (a=42)
            yield {0: DeltaToolCall(name=info.output_tools[0].name, json_args='{"a": 42')}
            yield {0: DeltaToolCall(json_args=', "b": "f')}
            yield {0: DeltaToolCall(json_args='oo"}')}

    agent = Agent(FunctionModel(stream_function=stream_model), output_type=Foo)

    @agent.output_validator
    def validate_output(ctx: RunContext[None], output: Foo) -> Foo:
        assert ctx.tool_name == 'final_result'
        # Only validate on final output, not partial streaming outputs
        if not ctx.partial_output:
            if output.a == 42:
                return output
            else:
                raise ModelRetry('"a" should be 42')
        return output

    async with agent.run_stream('Hello') as result:
        outputs = [output async for output in result.stream_output(debounce_by=None)]

    # Verify final output is correct after retry
    assert outputs[-1] == Foo(a=42, b='foo')
    assert result.output == Foo(a=42, b='foo')

    # Verify retry happened by checking message history
    messages = result.all_messages()
    assert len(messages) == 5  # Request, Response (wrong), Retry, Response (correct), ToolReturn
    assert isinstance(messages[2].parts[0], RetryPromptPart)
    assert messages[2].parts[0].content == '"a" should be 42'
    assert messages[2].parts[0].tool_name == 'final_result'

Python, Pydantic AI & LLM client version

python 3.10
pydantic-ai 1.11.0

[DouweM]

@petersli I agree this shouldn't blow up the way it does, but it's not trivial to fix.

The note from https://ai.pydantic.dev/agents/#streaming-events-and-final-output is relevant here:

As the run_stream() method will consider the first output matching the output type to be the final output, it will stop running the agent graph and will not execute any tool calls made by the model after this "final" output.

Basically, once run_stream() finds the (start of) final output in the response, it breaks out of the agent graph and yields the StreamedRunResult that you can then use to stream to the end of that response using methods like stream_output() and stream_text(). The expectation of those methods is that they incrementally stream output as new deltas come in, either as growing output objects or text, or as text deltas.

If ModelRetry is raised during output streaming, we'd effectively abandon the original output, go back to the the agent graph, send a new model request, and start getting deltas for a complete new object.

That's incompatible with how those methods currently (are expected to) work:

• stream_output() could start yielding the new object as it comes in, but that would break the expectation that "the yielded output object always grows" that some users rely on, for example when the output type is a list, and you keep track of the length of the yielded list to pull out completed list items -- that'll break if the list suddenly shrinks
• stream_output() has the same issue with delta=False, and with delta=True it's not clear what the right behavior would be at all.

When ModelRetry is used with stream_output(), breaking that "the yielded output object always grows" expectation is probably better than blowing up with an exception, but it's especially the stream_text(delta=True) case where I'm not sure what we should do. Perhaps it'd be worth fixing for stream_output() and stream_text(delta=False) and explicitly documenting that stream_text(delta=True) limitation.

Either way this would take a bit of refactoring to make work as run_stream currently breaks out of the agent graph entirely once it yields the StreamedRunResult which wraps only the StreamedResponse from the model, and we'd instead need it to wrap the entire agent graph...

[petersli]

@DouweM

I guess in all cases (stream_output and stream_text with delta=True and delta=False), I'd expect the ModelRetry to go back to the agent graph and send a new model request, and put the onus on the the user to make the error message in ModelRetry prevent any issues with repeated output during streaming.

E.g. in my implementation, I was doing something like:

raise ModelRetry(
                f"Invalid [field name], try again: {e.message}.
 Note: the text field has already been streamed to the user -- on your new response, `text` should be an empty string."
            )

In effect, I think anything that has already happened during streaming before the ModelRetry is raised should be assumed to be immutable / already streamed to the end user (including text deltas), and recovering properly should be a prompting concern.

• This is already how it works with ModelRetry in function tools if the model e.g. yields some text parts, then calls the tool and gets a ModelRetry

[DouweM]

@petersli What would that mean exactly for the behavior of stream_text(delta=True)? If the text doesn't pass validation, the model should output new text -- prompting it not do would defeat the point wouldn't it? Thinking of possible workarounds, I suppose we could yield some special sentinel value to indicates that the streamed text should be reset.

• This is already how it works with ModelRetry in function tools if the model e.g. yields some text parts, then calls the tool and gets a ModelRetry

The difference is that that only happens when you're streaming all events, right, not just output or text, so you get clear start/delta/end events and can know when original text should be discarded.

recovering properly should be a prompting concern.

I agree recovering properly should be the user's concern, whether through prompting or careful handling of what stream_{output,text}() yield. Since they'd currently get the ModelRetry bubbled up, it seems strictly better to make this work somehow, but I don't think we can get around giving stream_text(delta=True) some way to communicate that it started streaming new output, not deltas to append to what was already streamed.

[petersli]

@DouweM I just realized, don't we only run validators when delta=False?

pydantic-ai/pydantic_ai_slim/pydantic_ai/result.py

Line 102 in 53962e5

Result validators will NOT be called on the text result if `delta=True`.

so solving for stream_output and stream_text(delta=False) should be sufficient? unless I'm missing something

[DouweM]

@petersli Ha, that's a very good point! My bad for getting us on a wrong track.

I suppose I'm partially pushing back because I think this would be less than trivial to implement because the StreamedRunResult these methods exist on currently wraps a single StreamedResponse from the model, and we'd instead need it to wrap the entire agent graph and multiple streamed responses.

Would you be up for having a try at that by any chance? :)