Llama4 with ToolOutput using PydanticAI

pydanticai_bedrock_llama4_tooloutput.py

# Either AWS credentials (AWS_SECRET_KEY or AWS_ACCESS_KEY_ID must be set before hand)
# OR also, you can use AWS_BEARER_TOKEN_BEDROCK

from pydantic import BaseModel
from pydantic_ai import Agent, ToolOutput
from pydantic_ai.providers.bedrock import BedrockModelProfile, BedrockProvider
from pydantic_ai.models.bedrock import BedrockConverseModel

MODEL_NAME = 'us.meta.llama4-scout-17b-instruct-v1:0'

class Fruit(BaseModel):
    name: str
    color: str


class Vehicle(BaseModel):
    name: str
    wheels: int


model = BedrockConverseModel(
    model_name=MODEL_NAME,
    profile=BedrockModelProfile(bedrock_supports_tool_choice=False)
)

agent = Agent(
    model,
    output_type=[
        ToolOutput(Fruit, name='return_fruit'),
        ToolOutput(Vehicle, name='return_fruit')
    ]
)

result = await agent.run('What is a banana?')
result

result = await agent.run('What is a 4-wheeled ATV?')
result