from pydantic_ai import Agent, ModelRetry
from random import randint
import asyncio

def get_random_number():
    """
    Return a random number using the randint
    """
    res = randint(0, 100)
    if res > 50:
        raise ModelRetry("Error: Random number is greater than 50")
    return res

agent = Agent(
    "bedrock:us.anthropic.claude-3-7-sonnet-20250219-v1:0",
    system_prompt="You are a helpful assistant that generates random numbers.",
    retries=3,
    tools=[get_random_number],
)

async def main():
    async with agent.iter(user_prompt="Generate 10 different random numbers and then create a summary of these numbers.") as result:
        async for node in result:
            if Agent.is_model_request_node(node):
                async with node.stream(result.ctx) as request_stream:
                    async for event in request_stream:
                        print(event)

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