from __future__importannotationsfromdataclassesimportdataclassfromtyping_extensionsimport (Generic,Sequence,TypeVar,assert_type,)T=TypeVar("T")@dataclassclassAgent(Generic[T]):output_type:Sequence[type[T]]classFoo:passclassBar:pass# pyright - works# mypy - error: Expression is of type "Agent[object]", not "Agent[Foo | Bar]"  [assert-type]# pyrefly - assert_type(Agent[Foo], Agent[Bar | Foo]) failed + Argument `list[type[Bar] | type[Foo]]` is not assignable to parameter `output_type` with type `Sequence[type[Foo]]` in function `Agent.__init__`# ty - `Agent[Foo | Bar]` and `Agent[Unknown]` are not equivalent typesassert_type(Agent([Foo,Bar]),Agent[Foo|Bar])# pyright - works# mypy - error: Expression is of type "Agent[Never]", not "Agent[int | str]"  [assert-type]# pyrefly - assert_type(Agent[int], Agent[str | int]) failed + Argument `list[type[str] | type[int]]` is not assignable to parameter `output_type` with type `Sequence[type[int]]` in function `Agent.__init__`# ty - `Agent[int | str]` and `Agent[Unknown]` are not equivalent typesassert_type(Agent([int,str]),Agent[int|str])# worksassert_type(Agent[Foo|Bar]([Foo,Bar]),Agent[Foo|Bar])# worksassert_type(Agent[int|str]([int,str]),Agent[int|str])

fromdataclassesimportdataclassfromtyping_extensionsimport (Callable,Generic,TypeVar,assert_type,)T=TypeVar("T")@dataclassclassAgent(Generic[T]):output_type:Callable[...,T]deffunc()->int:return1# pyright, mypy, pyrefly - works# ty - `Agent[int]` and `Agent[Unknown]` are not equivalent types + Expected `((...) -> T) | ((...) -> Awaitable[T])`, found `def func() -> int`assert_type(Agent(func),Agent[int])# worksassert_type(Agent[int](func),Agent[int])

fromdataclassesimportdataclassfromtyping_extensionsimport (Awaitable,Callable,Generic,TypeVar,assert_type,)T=TypeVar("T")@dataclassclassAgent(Generic[T]):output_type:Callable[...,T]|Callable[...,Awaitable[T]]asyncdefcoro()->bool:returnTruedeffunc()->int:return1# pyright, mypy, pyrefly - works# ty - `Agent[int]` and `Agent[Unknown]` are not equivalent types + Expected `((...) -> T) | ((...) -> Awaitable[T])`, found `def func() -> int`assert_type(Agent(func),Agent[int])# mypy - error: Argument 1 to "Agent" has incompatible type "Callable[[], Coroutine[Any, Any, bool]]"; expected "Callable[..., Never] | Callable[..., Awaitable[Never]]"  [arg-type]coro_agent=Agent(coro)# pyright, pyrefly - works# mypy - error: Expression is of type "Agent[Any]", not "Agent[bool]"# ty - `Agent[bool]` and `Agent[Unknown]` are not equivalent typesassert_type(coro_agent,Agent[bool])# worksassert_type(Agent[bool](coro),Agent[bool])

@dataclassclassOutput(Generic[T]):output_type:type[T]defor_[S](self,output_type:type[S])->Output[T|S]:returnOutput(self.output_type|output_type)# type: ignore# ordefoutput_type[T](*args:type[T])->Output[T]:raiseNotImplementedError

fromdotenvimportload_dotenvfrompydantic_aiimportAgent,RunContextload_dotenv()maths_agent=Agent(model="google-gla:gemini-2.0-flash",instructions="You are a maths tutor. Given a question, you will provide a step by step solution.",)asyncdefhand_off_to_maths_agent(ctx:RunContext,query:str)->str:res=awaitmaths_agent.run(query)ctx.messages+=res.new_messages()returnres.outputpoet_agent=Agent(model="google-gla:gemini-2.0-flash",instructions="You are a poet. Given a topic, you will provide a poem.",)asyncdefhand_off_to_poet_agent(ctx:RunContext,query:str)->str:res=awaitpoet_agent.run(query)ctx.messages+=res.new_messages()returnres.outputrouter_agent=Agent(model="google-gla:gemini-2.0-flash",instructions="You are a router. Given a user query, you will route it to the appropriate agent.",output_type=[hand_off_to_maths_agent,hand_off_to_poet_agent],)asyncdefmain():query="Calculate 10 + 10"result=awaitrouter_agent.run(query)formessageinresult.all_messages():print(message,"\n")print(result.output)if__name__=="__main__":importasyncioasyncio.run(main())

...async def hand_off_to_maths_agent(ctx: RunContext, query: str) -> str:    res = await maths_agent.run(query)    ctx.messages += res.new_messages()    return res.output...async def main():    query = "Calculate 10 + 10"    result = await router_agent.run(query)        query = "Write a poem about your answer"    result = await router_agent.run(query, message_history=result.all_messages())

pydantic_ai.exceptions.ModelHTTPError: status_code: 400, model_name: gpt-4o, body: {'message': "An assistant message with 'tool_calls' must be followed by tool messages responding to each 'tool_call_id'. The following tool_call_ids did not have response messages: call_KNSIrALYz0K6cTvpEIVVV7lW", 'type': 'invalid_request_error', 'param': 'messages.[3].role', 'code': None}

ModelRequest(parts=[UserPromptPart(content='Calculate 10 + 10', timestamp=datetime.datetime(2025, 6, 4, 0, 33, 24, 203261, tzinfo=datetime.timezone.utc))], instructions='You are a router. Given a user query, you will route it to the appropriate agent.') ModelResponse(parts=[ToolCallPart(tool_name='final_result_hand_off_to_maths_agent', args='{"query":"Calculate 10 + 10"}', tool_call_id='call_KNSIrALYz0K6cTvpEIVVV7lW')], usage=Usage(requests=1, request_tokens=118, response_tokens=25, total_tokens=143, details={'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0, 'cached_tokens': 0}), model_name='gpt-4o-2024-08-06', timestamp=datetime.datetime(2025, 6, 4, 0, 33, 24, tzinfo=datetime.timezone.utc), vendor_id='chatcmpl-BeWMGkev6eeWK9arEZ17o3OVyWUan') ModelRequest(parts=[UserPromptPart(content='Calculate 10 + 10', timestamp=datetime.datetime(2025, 6, 4, 0, 33, 25, 554116, tzinfo=datetime.timezone.utc))], instructions='You are a maths tutor. Given a question, you will provide a step by step solution.') ModelResponse(parts=[TextPart(content='To calculate \\(10 + 10\\), simply add the two numbers together:\n\n\\[ \n10 + 10 = 20 \n\\]\n\nSo, the answer is \\(20\\).')], usage=Usage(requests=1, request_tokens=46, response_tokens=37, total_tokens=83, details={'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0, 'cached_tokens': 0}), model_name='gpt-4o-2024-08-06', timestamp=datetime.datetime(2025, 6, 4, 0, 33, 25, tzinfo=datetime.timezone.utc), vendor_id='chatcmpl-BeWMHmqsk9Nip3TEfkAbHZevvLe4V') ModelRequest(parts=[ToolReturnPart(tool_name='final_result_hand_off_to_maths_agent', content='Final result processed.', tool_call_id='call_KNSIrALYz0K6cTvpEIVVV7lW', timestamp=datetime.datetime(2025, 6, 4, 0, 33, 26, 305492, tzinfo=datetime.timezone.utc))])

fromdataclassesimportreplacefromdotenvimportload_dotenvfrompydantic_aiimportAgent,RunContextfrompydantic_ai.messagesimportModelMessage,ModelRequest,ToolCallPart,ToolReturnPartload_dotenv()maths_agent=Agent(model="google-gla:gemini-2.0-flash",instructions="You are a maths tutor. Given a question, you will provide a step by step solution.",)asyncdefhand_off_to_maths_agent(ctx:RunContext,query:str)->str:res=awaitmaths_agent.run(query)ctx.messages+=res.new_messages()returnres.outputpoet_agent=Agent(model="google-gla:gemini-2.0-flash",instructions="You are a poet. Given a topic, you will provide a poem.",)asyncdefhand_off_to_poet_agent(ctx:RunContext,query:str)->str:res=awaitpoet_agent.run(query)ctx.messages+=res.new_messages()returnres.outputrouter_agent=Agent(model="google-gla:gemini-2.0-flash",instructions="You are a router. Given a user query, you will route it to the appropriate agent.",output_type=[hand_off_to_maths_agent,hand_off_to_poet_agent],)deffilter_tool_parts(messages:list[ModelMessage],filter_str:str)->list[ModelMessage]:filtered_messages:list[ModelMessage]= []formessageinmessages:ifisinstance(message,ModelRequest):filtered_parts= [partforpartinmessage.partsifnot (isinstance(part,ToolReturnPart)andfilter_strinpart.tool_name)            ]iffiltered_parts:filtered_messages.append(replace(message,parts=filtered_parts))else:filtered_parts= [partforpartinmessage.partsifnot (isinstance(part,ToolCallPart)andfilter_strinpart.tool_name)            ]iffiltered_parts:filtered_messages.append(replace(message,parts=filtered_parts))returnfiltered_messagesasyncdefmain():query="Calculate 10 + 10"result=awaitrouter_agent.run(query)message_history=filter_tool_parts(result.all_messages(),"hand_off")sep="\n"+"-"*100+"\n"print(sep)formessageinmessage_history:print(message,"\n")print(f"{sep}{result.output}{sep}")query="Write a poem about your answer"result=awaitrouter_agent.run(query,message_history=message_history)print(result.output)if__name__=="__main__":importasyncioasyncio.run(main())