- Notifications
You must be signed in to change notification settings - Fork1k
Description
Initial Checks
- I confirm that I'm using the latest version of Pydantic AI
- I confirm that I searched for my issue inhttps://github.com/pydantic/pydantic-ai/issues before opening this issue
Description
This is a follow-up from a Slack conversation:https://pydanticlogfire.slack.com/archives/C083V7PMHHA/p1751066192272409
I want my agent to stream text, make tool calls, stream more text if needed, and make a final tool call for its structured output (I'm using Gemini which does not support tool calls and native structured output at the same time). To do so, I have set output_type to[str, FinalOutput] where FinalOutput is my structured output type.
This streams text as expected, but the final_result tool call is not always made. My understanding is that since the end node can either be FinalOutput or text, it will sometimes output text and there is no guarantee it outputs a FinalOutput. In my repro script, I have two example queries. One of them will call the final output tool call, and one of them will not (although this is dependent on the model's output being reproducible which isn't guaranteed).
I think what I would ideally want to do is to set output_type toFinalOutput so that it is always called but still enable streaming text output. That is currently prevented bythis line in _agent_graph.py.
Is there a way to have pydantic-ai consistently stream text output AND make a final result tool call? One workaround is to just prompt the LLM to make the final result tool call, but that is not guaranteed.
Example Code
"""To run: python <name>.py.Requires GEMINI_API_KEY environment variable (available at https://aistudio.google.com/apikey).Alternatively you can change the model but this will change the reproducibility of the script."""frompydanticimportBaseModelfrompydantic_aiimportAgentfrompydantic_ai.messagesimport (FinalResultEvent,FunctionToolCallEvent,FunctionToolResultEvent,PartDeltaEvent,PartStartEvent,TextPart,TextPartDelta,)frompydantic_ai.modelsimportKnownModelNameclassFinalOutput(BaseModel):lines:list[str]QUERY="Write a short poem about Python."# Does not call final tool call.# QUERY = "Write a poem about Python." # Does call final tool call.MODEL:KnownModelName="google-gla:gemini-2.5-flash"AGENT=Agent(model=MODEL,model_settings={"temperature":0},output_type=[str,FinalOutput],)asyncdefmain():# Based on https://ai.pydantic.dev/agents/#streamingasyncwithAGENT.iter(user_prompt=QUERY)asrun:asyncfornodeinrun:ifAgent.is_user_prompt_node(node):print("\nUser prompt:",node.user_prompt)elifAgent.is_model_request_node(node):asyncwithnode.stream(run.ctx)asrequest_stream:asyncforeventinrequest_stream:ifisinstance(event,PartStartEvent)andisinstance(event.part,TextPart):print("\nText start:",event.part.content)elifisinstance(event,PartDeltaEvent)andisinstance(event.delta,TextPartDelta):print("\nText delta:",event.delta.content_delta)elifisinstance(event,FinalResultEvent):print("\nFinal output tool call:",event.tool_name)elifAgent.is_call_tools_node(node):asyncwithnode.stream(run.ctx)ashandle_stream:asyncforeventinhandle_stream:ifisinstance(event,FunctionToolCallEvent):print("\nTool call event:",event)elifisinstance(event,FunctionToolResultEvent):print("\nTool result event:",event)elifAgent.is_end_node(node):print("\nFinal output:",node.data.output)if__name__=="__main__":importasyncioasyncio.run(main())
Python, Pydantic AI & LLM client version
Python 3.11.13pydantic-ai 0.4.0