State in Custom Function

This example demonstrates how to access the run context in a custom python function step

1. Create a Python file

   Create a file named access_session_state_in_custom_python_function_step.py

2. Add code to file

   access_session_state_in_custom_python_function_step.py

   from agno.agent import Agent
   from agno.db.sqlite import SqliteDb
   from agno.models.openai import OpenAIResponses
   from agno.run import RunContext
   from agno.team import Team
   from agno.tools.hackernews import HackerNewsTools
   from agno.tools.yfinance import YFinanceTools
   from agno.workflow.step import Step, StepInput, StepOutput
   from agno.workflow.workflow import Workflow
   
   # Define agents
   hackernews_agent = Agent(
       name="Hackernews Agent",
       model=OpenAIResponses(id="gpt-5.2"),
       tools=[HackerNewsTools()],
       instructions="Extract key insights and content from Hackernews posts",
   )
   
   finance_agent = Agent(
       name="Finance Agent",
       model=OpenAIResponses(id="gpt-5.2"),
       tools=[YFinanceTools()],
       instructions="Get financial data and market trends",
   )
   
   # Define research team for complex analysis
   research_team = Team(
       name="Research Team",
       model=OpenAIResponses(id="gpt-5.2"),
       members=[hackernews_agent, finance_agent],
       instructions="Analyze content and create comprehensive social media strategy",
   )
   
   content_planner = Agent(
       name="Content Planner",
       model=OpenAIResponses(id="gpt-5.2"),
       instructions=[
           "Plan a content schedule over 4 weeks for the provided topic and research content",
           "Ensure that I have posts for 3 posts per week",
       ],
   )
   
   
   def custom_content_planning_function(
       step_input: StepInput, run_context: RunContext
   ) -> StepOutput:
       """
       Custom function that does intelligent content planning with context awareness
       and maintains a content plan history in session_state
       """
       message = step_input.input
       previous_step_content = step_input.previous_step_content
   
       # Initialize content history if not present
       if "content_plans" not in run_context.session_state:
           run_context.session_state["content_plans"] = []
   
       if "plan_counter" not in run_context.session_state:
           run_context.session_state["plan_counter"] = 0
   
       # Increment plan counter
       run_context.session_state["plan_counter"] += 1
       current_plan_id = run_context.session_state["plan_counter"]
   
       # Create intelligent planning prompt
       planning_prompt = f"""
           STRATEGIC CONTENT PLANNING REQUEST:
   
           Core Topic: {message}
           Plan ID: #{current_plan_id}
   
           Research Results: {previous_step_content[:500] if previous_step_content else "No research results"}
   
           Previous Plans Count: {len(run_context.session_state["content_plans"])}
   
           Planning Requirements:
           1. Create a comprehensive content strategy based on the research
           2. Leverage the research findings effectively
           3. Identify content formats and channels
           4. Provide timeline and priority recommendations
           5. Include engagement and distribution strategies
   
           Please create a detailed, actionable content plan.
       """
   
       try:
           response = content_planner.run(planning_prompt)
   
           # Store this plan in session state
           plan_data = {
               "id": current_plan_id,
               "topic": message,
               "content": response.content,
               "timestamp": f"Plan #{current_plan_id}",
               "has_research": bool(previous_step_content),
           }
           run_context.session_state["content_plans"].append(plan_data)
   
           enhanced_content = f"""
               ## Strategic Content Plan #{current_plan_id}
   
               **Planning Topic:** {message}
   
               **Research Integration:** {"Research-based" if previous_step_content else "No research foundation"}
               **Total Plans Created:** {len(run_context.session_state["content_plans"])}
   
               **Content Strategy:**
               {response.content}
   
               **Custom Planning Enhancements:**
               - Research Integration: {"High" if previous_step_content else "Baseline"}
               - Strategic Alignment: Optimized for multi-channel distribution
               - Execution Ready: Detailed action items included
               - Session History: {len(run_context.session_state["content_plans"])} plans stored
   
               **Plan ID:** #{current_plan_id}
           """.strip()
   
           return StepOutput(content=enhanced_content)
   
       except Exception as e:
           return StepOutput(
               content=f"Custom content planning failed: {str(e)}",
               success=False,
           )
   
   
   def content_summary_function(step_input: StepInput, run_context: RunContext) -> StepOutput:
       """
       Custom function that summarizes all content plans created in the session
       """
       if run_context.session_state is None or run_context.session_state.get("content_plans") is None:
           return StepOutput(
               content="No content plans found in session state.", success=False
           )
   
       plans = run_context.session_state["content_plans"]
       summary = f"""
           ## Content Planning Session Summary
   
           **Total Plans Created:** {len(plans)}
           **Session Statistics:**
           - Plans with research: {len([p for p in plans if p["has_research"]])}
           - Plans without research: {len([p for p in plans if not p["has_research"]])}
   
           **Plan Overview:**
       """
   
       for plan in plans:
           summary += f"""
   
           ### Plan #{plan["id"]} - {plan["topic"]}
           - Research Available: {"Yes" if plan["has_research"] else "No"}
           - Status: Completed
           """
   
       # Update session state with summary info
       run_context.session_state["session_summarized"] = True
       run_context.session_state["total_plans_summarized"] = len(plans)
   
       return StepOutput(content=summary.strip())
   
   
   # Define steps using different executor types
   
   research_step = Step(
       name="Research Step",
       team=research_team,
   )
   
   content_planning_step = Step(
       name="Content Planning Step",
       executor=custom_content_planning_function,
   )
   
   content_summary_step = Step(
       name="Content Summary Step",
       executor=content_summary_function,
   )
   
   
   # Define and use examples
   if __name__ == "__main__":
       content_creation_workflow = Workflow(
           name="Content Creation Workflow",
           description="Automated content creation with custom execution options and session state",
           db=SqliteDb(
               session_table="workflow_session",
               db_file="tmp/workflow.db",
           ),
           # Define the sequence of steps
           # First run the research_step, then the content_planning_step, then the summary_step
           # You can mix and match agents, teams, and even regular python functions directly as steps
           steps=[research_step, content_planning_step, content_summary_step],
           # Initialize session state with empty content plans
           session_state={"content_plans": [], "plan_counter": 0},
       )
   
       print("=== First Workflow Run ===")
       content_creation_workflow.print_response(
           input="AI trends in 2024",
           markdown=True,
       )
   
       print(
           f"\nSession State After First Run: {content_creation_workflow.get_session_state()}"
       )
   
       print("\n" + "=" * 60 + "\n")
   
       print("=== Second Workflow Run (Same Session) ===")
       content_creation_workflow.print_response(
           input="Machine Learning automation tools",
           markdown=True,
       )
   
       print(f"\nFinal Session State: {content_creation_workflow.get_session_state()}")
   
3. Set up your virtual environment

   Mac

   uv venv --python 3.12
   source .venv/bin/activate
   

   uv venv --python 3.12
   .venv\Scripts\activate
   
4. Install dependencies

   uv pip install -U agno openai yfinance
   
5. Export your OpenAI API key

   Set OpenAI Key

   Set your OPENAI_API_KEY as an environment variable. You can get one from OpenAI.

   Mac

   export OPENAI_API_KEY=sk-***
   

   setx OPENAI_API_KEY sk-***
   
6. Run Workflow

   python access_session_state_in_custom_python_function_step.py