How to stream from inside a tool?

[pedroallenrevez]

Imagine the following case:

from pydantic_ai import Agent, RunContext

qa_agent = Agent(  
    'openai:gpt-4o',
    system_prompt=(
        'Use the `qa` function to reply to user's questions about pydantic ai '
    ),
)

@qa.tool
async def qa(ctx: RunContext[int], query: str) -> AsyncGenerator[str | Reference, None]:  
    """Answer users questions about pydantic ai, by inspecting documentation"""
    docs = get_context(...)
    async for token in prompt(docs, ...):
        yield token # Can be a Reference to a document, or just text 

This would also need the agent to end earlier, which is a known issue that other people are asking for.
I also know that this can be circumvented by using the graph implementation.
But I also don't see a way to stream results from inside of a node in the graph implementation.
Is this a pattern that is accounted for, or is planned for the future?
Can we do this right now?

[pedroallenrevez]

@Kludex ? Any idea?

[Kludex]

Hi @pedroallenrevez , sorry the delay.

Can you try to explain a use case that you'd need this? I'm having a hard time seeing a real case scenario that you'd want a generator.

[wylansford]

@Kludex

This is quite useful for frontend functionality. One specific example is tracking the progress of a longer running tool call and streaming those results in real time to the user.

[A3Ackerman]

@wylansford I'm achieving something similar by passing a buffer/queue into tools via Deps. All my responses are written to the buffer instead of directly yielded to the caller.