Stream a function step through AgentOS

from typing import AsyncIterator, Union

from agno.agent.agent import Agent
from agno.db.in_memory import InMemoryDb
from agno.db.sqlite import SqliteDb
from agno.models.openai import OpenAIChat
from agno.os import AgentOS
from agno.run.agent import RunOutputEvent
from agno.tools.hackernews import HackerNewsTools
from agno.workflow.step import Step, StepInput, StepOutput
from agno.workflow.workflow import Workflow

research_agent = Agent(
    name="Hacker News Researcher",
    model=OpenAIChat(id="gpt-4o"),
    tools=[HackerNewsTools()],
    instructions="Find relevant Hacker News stories and summarize the key points.",
)

content_planner = Agent(
    name="Content Planner",
    model=OpenAIChat(id="gpt-4o"),
    instructions="Create an actionable content plan from the supplied research.",
    db=InMemoryDb(),
)


async def stream_content_plan(
    step_input: StepInput,
) -> AsyncIterator[Union[RunOutputEvent, StepOutput]]:
    """Stream planner events, then return the completed content plan."""
    message = step_input.input
    previous_step_content = step_input.previous_step_content

    planning_prompt = f"""
    Topic: {message}
    Research: {previous_step_content or "No research results"}

    Create a four-week content plan with three posts per week.
    """

    try:
        events = content_planner.arun(
            planning_prompt, stream=True, stream_events=True
        )
        async for event in events:
            yield event

        response = content_planner.get_last_run_output()
        yield StepOutput(content=response.content if response else "")
    except Exception as exc:
        yield StepOutput(
            content=f"Content planning failed: {exc}",
            success=False,
        )


research_step = Step(
    name="Research Step",
    agent=research_agent,
)

content_planning_step = Step(
    name="Content Planning Step",
    executor=stream_content_plan,
)

streaming_content_workflow = Workflow(
    name="Streaming Content Creation Workflow",
    description="Automated content creation with streaming custom execution functions",
    db=SqliteDb(
        session_table="workflow_session",
        db_file="tmp/workflow.db",
    ),
    steps=[
        research_step,
        content_planning_step,
    ],
)

agent_os = AgentOS(
    workflows=[streaming_content_workflow],
)
app = agent_os.get_app()

if __name__ == "__main__":
    agent_os.serve(app="step_with_function_streaming_agentos:app", reload=True)