Do `output_type` functions support streaming input?

[Danipulok]

Question

Start

Hey here!
Thank you very much for such a good library!

I've read the docs, but didn't seem to find any info on streaming in output_types function?
Is it supported, or what's the intended method to stream output result inside functions?

Our use case:

We always use output types (we call them output functions) to stream to UI

Example:

class CarouselItem(BaseModel):
    text: str | None = None
    url: HttpUrl | None = None

class CarouselMessage(BaseModel):
    text: str | None = None  # to enable streaming from the start
    items: list[CarouselItem] | None = None  # to enable streaming from the start

async def send_carousel_message(message: CarouselMessage, ctx: Ctx) -> None: ...
    await ctx.ui.send_message(message)

Is there a way to do similar but with PydanticAI?

MRE

Here's some MRE I wrote to show what I would like to get:

import asyncio

from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider


class CityLocation(BaseModel):
    city: str | None = None
    country: str | None = None


async def show_message(
    data: CityLocation,
) -> None:
    """Display the city location information."""
    print(f"City: {data.city}, Country: {data.country}")  # currently there's only one print and result is not streamed to here


async def main() -> None:
    """Run the agent."""
    api_key = "..."
    model = OpenAIChatModel(
        "gpt-4o",
        provider=OpenAIProvider(api_key=api_key),
    )
    agent = Agent(model, output_type=show_message)
    await agent.run("Where were the olympics held in 2012?")


if __name__ == "__main__":
    asyncio.run(main())

Motivation

IMHO it would be really useful to incapsulate this logic inside functions, and do not have something like:

output_objects = list[Function1, Function2, Function3]
output_type = Union[*output_objects]
async with agent.run_stream(...) as result:
    output_message = TypeAdapter(output_type).movel_validate(result.output)
    await channel.send(output_message)

Additional Context

No response

[DouweM]

@Danipulok Is #3114 what you're looking for? That's primarily built for tool functions, but will work for output functions as well.

[Danipulok]

@DouweM, hi, sorry for the late response!
And thank you for replying

Is #3114 what you're looking for?

I guess not really, but that idea seems useful too

I was a little bit wrong in my MRE earlier, because didn't want to provide full example to burden you, I guess :D
And thought maybe I skipped something in the docs
Sorry about that

So our tool protocol is following:

class BaseTool[InputT, OutputT](abc.ABC):
    @abc.abstractmethod
    async def call(self, input_: InputT) -> OutputT: ...

    async def stream(
        self,
        input_stream: AsyncIterator[InputT],
    ) -> AsyncIterator[OutputT]:
        """Stream input through the tool.
        By default, this method calls the tool **only** on the latest input item.
        """
        _chunk: InputT | None = None
        async for _chunk in input_stream:
            ...

        if _chunk is None:
            msg = "Input stream is empty"
            raise ValueError(msg)

        yield await self.call(_chunk)

And real tools can define only call, if they are "sync" (calculations, DB queries, or something that does not have to stream the input to other services)

Or they can define stream too, if they are used for streaming LLM responses to user, for example (I can't really think of any other examples)
Then it would look like this:

class SendMessageTool(BaseTool[Message, None]):
    async def call(self, input_: Message) -> None:
        await channel.send_message(input_)  # user has to wait until the whole message is received

    async def stream(self, input_: AsyncIterator[Message]) -> None:
        async for chunk in input_:
            await channel.send_delta(chunk)  # user sees streamed message, as everywhere this days

Basically in your API it would accept what agent.stream_responses returns (AsyncIterator[tuple[_messages.ModelResponse, bool]]), but parsed into structures (or at least raw input, maybe?)

So the issue question is do your support streaming responses from LLMs inside tools or what's the intended way to do it?

[Danipulok]

I created a new issue #3485 with a different question, there is an example of how we would use Pydantic AI with streaming and in real-world scenario
Each output message type is a separate tool to increase probability of a better LLM response

[DouweM]

@Danipulok Aah now I think I understand better: you're thinking of a tool/function that takes an LLM response stream as input? And it sounds like you want the LLM to choose which tool/function to call?

If that's right:

• Would you want the function to take ModelResponses?
• Or just output tool/function args, similar to what you'd get from run_stream + stream_output?
• Or some kind of deltas?
• Or agent stream events?

[Danipulok]

@DouweM , hey again!
Sorry for the late response, I was consulting with my team members

you're thinking of a tool/function that takes an LLM response stream as input? And it sounds like you want the LLM to choose which tool/function to call?

Yes, you are right. To get:

1. llm input as tool stream input
2. in the used output tool

Would you want the function to

I think there are 3 options IMHO:

1. AsyncIterable[ModelResponses] (your option 1)
   Then is seems pretty customizable and user can do whatever he wants with it.
   Validate, create a new stream, pass it further

2. AsyncIterable[MyInputModel] (your option 2)
   This seems sufficient for our use-case and the easiest one.
   Then we can simply pass this stream everywhere we need

3. Combo of option1 and option2
   Maybe some new type from pydantic-ai internals with generic dataclass, that has parsed MyInputModel and raw models responses?
   But I'm not sure about it at all

Something like this, maybe:

@dataclass
class IterableInput[T: type[InputModelT]]:
    parsed: T
    delta: ModelResponse  # maybe list, not sure
    is_last_chunk: bool

Then the tool input is AsyncIterable[IterableInput[MyModel]]

Or some kind of deltas?

At least for us it's not our use-case. We accept accumulated response, and not deltas

Or agent stream events?

Not sure about it, but seems not the case either

Summary:

I think both option 1 and option 2 are good.
But option 1 seems more customizable and maybe it's best to use it, and not just accept parsed input.
Or maybe there's a way to do the best of those two

[DouweM]

@Danipulok 2 (or 3) fit better in the Pydantic AI model, as we use the output function arguments type hints to determine the JSON schema the model should generate. If the function is just hinted as AsyncIterator[ModelResponse], we wouldn't know what shape the model output should have.

In any case, I think this'd be an interesting feature but it's not trivial to implement and pretty niche, so it's not something I'd expect us to prioritize any time soon, if ever. If you look into contributing it though, I'd be happy to give some hints if you get stuck!

[Danipulok]

@DouweM, hey here!
I discussed with my team, I think I will be doing a PR to implement this

Are there are directions you could give me right now?

Here are some questions I have right now:

1. It seems I will have to add takes_streaming: bool = False to the class Tool(Generic[ToolAgentDepsT]): ?
2. Will you accept the PR if I implement streaming support only for output tools, and not 'usual' tools?
3. I think it should look something like this:

class InputModel(BaseModel):
    foo: str | None = None
    bar: str | None = None

agent = Agent()

@agent.plain_tool(taskes_streaming=True)
async def async_streaming_tool(stream: AsyncIterable[InputModel]) -> None:
    async for chunk in stream:
        print(chunk)  
#> "Hello"
#> "Hello, World!"


@agent.plain_tool(taskes_streaming=True)
def sync_streaming_tool(stream: Iterable[InputModel]) -> None:
    for chunk in stream:
        print(chunk)
#> "Hello"
#> "Hello, World!"

async def ouput_async_streaming_tool(stream: AsyncIterable[InputModel]) -> None:
    async for chunk in stream:
        print(chunk)
#> "Hello"
#> "Hello, World!"

def output_sync_streaming_tool(stream: Iterable[InputModel]) -> None:
    for chunk in stream:
        print(chunk)
#> "Hello"
#> "Hello, World!"

output_type = Union[
    ToolOutput(ouput_async_streaming_tool, takes_streaming=True),
    ToolOutput(output_sync_streaming_tool, takes_streaming=True),
]
agent._output_type = output_type

Does it seems okay or should I change the API?

[DouweM]

@Danipulok I'd prefer to implement this only for output functions (both when used with output tools or with native output), not regular tools. I don't think we need a takes_streaming argument, we can detect if the signature is an AsyncIterable[T] and base it on that.

[Danipulok]

@DouweM, got it, thank you!
Is there anything else, or I can start working on the PR?

[DouweM]

@Danipulok Feel free to start working on it!