from pydantic_ai import Agent
from pydantic_ai.durable_exec.prefect import PrefectAgent

# Existing agent
agent = Agent('gpt-4o', name='my_agent')

# Enable durable execution
prefect_agent = PrefectAgent(agent)

# Use as normal, now run as a Prefect flow with model call and tool call tasks
result = await prefect_agent.run('What is the capital of Mexico?')

print(result.output) # The capital of Mexico is Mexico City

prefect_agent = PrefectAgent(
    agent,
    model_task_config=TaskConfig(
        retries=3,
        retry_delay_seconds=[1.0, 2.0, 4.0],
        timeout_seconds=30.0,
    ),
)

# Deploy agent with scheduled execution
from prefect import flow
from prefect.deployments import run_deployment
from pydantic_ai import Agent
from pydantic_ai.durable_exec.prefect import PrefectAgent

agent = Agent(
    'openai:gpt-4o',
    name='daily_report_agent',
    instructions='Generate a daily summary report.',
)

prefect_agent = PrefectAgent(agent)

@flow
async def daily_report_flow(user_prompt: str):
    """Generate a daily report using the agent."""
    result = await prefect_agent.run(user_prompt)
    return result.output

# Serve the flow with a daily schedule
if __name__ == '__main__':
    daily_report_flow.serve(
        name='daily-report-deployment',
        cron='0 9 * * *',  # Run daily at 9am
        parameters={'user_prompt': "Generate today's report"},
        tags=['production', 'reports'],
    )