Async MySQL for Workflow

import asyncio
import uuid
from typing import List

from agno.agent import Agent
from agno.db.base import SessionType
from agno.db.mysql import AsyncMySQLDb
from agno.tools.websearch import WebSearchTools
from agno.workflow.types import WorkflowExecutionInput
from agno.workflow.workflow import Workflow
from pydantic import BaseModel

db_url = "mysql+asyncmy://ai:ai@localhost:3306/ai"
db = AsyncMySQLDb(db_url=db_url)


class ResearchTopic(BaseModel):
    topic: str
    key_points: List[str]
    summary: str


researcher = Agent(
    name="Researcher",
    tools=[WebSearchTools()],
    instructions="Research the topic and return key points and a summary",
    output_schema=ResearchTopic,
)

writer = Agent(
    name="Writer",
    instructions="Write a blog post based on the research provided",
)


async def blog_workflow(workflow: Workflow, execution_input: WorkflowExecutionInput):
    topic = execution_input.input

    research_result = await researcher.arun(f"Research this topic: {topic}")

    if research_result and research_result.content:
        blog_result = await writer.arun(
            f"Write a blog post about {topic}. Use this research: {research_result.content.model_dump_json()}"
        )
        return blog_result.content

    return "Failed to complete workflow"


workflow = Workflow(
    name="Blog Generator",
    steps=blog_workflow,
    db=db,
)


async def main():
    try:
        session_id = str(uuid.uuid4())
        await workflow.aprint_response(
            input="The future of artificial intelligence",
            session_id=session_id,
            markdown=True,
        )
        session_data = await db.get_session(
            session_id=session_id, session_type=SessionType.WORKFLOW
        )
        print("\n=== SESSION DATA ===")
        print(session_data.to_dict())
    finally:
        await db.close()


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