Make it possible to stream structured while using .iter

[dmontagu]

Addresses the request here to support structured streaming while using .iter.

In particular, the analogous example to the whales streaming stuff would look like:

async with agent.iter(my_prompt) as run:
  async for node in run:
     ...
     # probably only want to do it one these nodes?
     elif Agent.is_model_request_node(node):
        async with node.stream(run.ctx) as request_stream:
          async for message, last in request_stream.stream_responses():
          maybe_whales = await run.validate_structured_output(message, allow_partial=not last)
          if maybe_whales is not None:
              whales = maybe_whales.output
    ...

I have some concerns about the type names etc. but I think overall this seems like a reasonable thing to add.

[dmontagu]

this error message and the one below just seemed confusing; I suspect they are hard/impossible to hit in practice but they should still include the relevant information imo

[dmontagu]

My biggest concern with this implementation is that it is annotated as returning a FinalResult. That type is exactly what you want to be returning here, I think (well, it would be enough to just use Option, but that's not public, and why not include tool info?). But the name FinalResult is kind of confusing when the purpose of adding this was to support cases where it literally is not the final result, i.e. partial streaming. On the other hand, it's pretty clear what's going on imo, and do we really want to introduce another Option-like result marker type? I'm open to others' opinions.

[dmontagu]

Not sure if we really want to add the complexity associated with allowing the end user to specify the tool, or just always look up the tool here. I'd be okay either way, dropping this should be trivial though if we want to (just delete the argument and the first if-statement referencing it below).

[DouweM]

This returns the same tuple[call, output_tool] we get from find_named_tool below, so if we're in this branch, we can skip that extra lookup

[github-actions]

Docs Preview

commit:	3ebc6a2
Preview URL:	https://c77af922-pydantic-ai-previews.pydantic.workers.dev

[DouweM]

This is basically a copy-paste of the other method in _agent_graph, I suggest moving it to a shared util.

[DouweM]

AgentStream already has a stream_output() method that works this way as of #2134.

"""Information about whales — an example of streamed structured response validation.

This script streams structured responses from GPT-4 about whales, validates the data
and displays it as a dynamic table using Rich as the data is received.

Run with:

    uv run -m pydantic_ai_examples.stream_whales
"""

from typing import Annotated

import logfire
from pydantic import Field
from rich.console import Console
from rich.live import Live
from rich.table import Table
from typing_extensions import NotRequired, TypedDict

from pydantic_ai import Agent
from pydantic_ai.output import NativeOutput

# 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
logfire.configure(send_to_logfire='if-token-present')
logfire.instrument_pydantic_ai()


class Whale(TypedDict):
    name: str
    length: Annotated[float, Field(description='Average length of an adult whale in meters.')]
    weight: NotRequired[
        Annotated[
            float,
            Field(description='Average weight of an adult whale in kilograms.', ge=50),
        ]
    ]
    ocean: NotRequired[str]
    description: NotRequired[Annotated[str, Field(description='Short Description')]]


agent = Agent('openai:gpt-4o', output_type=NativeOutput(list[Whale]))


async def main():
    console = Console()
    with Live('\n' * 36, console=console) as live:
        console.print('Requesting data...', style='cyan')
        async with agent.iter('Generate me details of 5 species of Whale.') as run:
            async for node in run:
                if Agent.is_model_request_node(node):
                    async with node.stream(run.ctx) as request_stream:
                        console.print('Response:', style='green')

                        async for whales in request_stream.stream_output(debounce_by=0.01):
                            table = Table(
                                title='Species of Whale',
                                caption='Streaming Structured responses from GPT-4',
                                width=120,
                            )
                            table.add_column('ID', justify='right')
                            table.add_column('Name')
                            table.add_column('Avg. Length (m)', justify='right')
                            table.add_column('Avg. Weight (kg)', justify='right')
                            table.add_column('Ocean')
                            table.add_column('Description', justify='right')

                            for wid, whale in enumerate(whales, start=1):
                                table.add_row(
                                    str(wid),
                                    whale['name'],
                                    f'{whale["length"]:0.0f}',
                                    f'{w:0.0f}' if (w := whale.get('weight')) else '…',
                                    whale.get('ocean') or '…',
                                    whale.get('description') or '…',
                                )
                            live.update(table)


if __name__ == '__main__':
    import asyncio

    asyncio.run(main())