00-playground/54-event-driven-agentic-document-workflows/L2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "def18024-676f-4733-b672-035365a0fcad",
   "metadata": {},
   "source": [
    "# Lesson 2: Building a Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64bad755-8d8c-4392-bcc0-dd7878314cee",
   "metadata": {},
   "source": [
    "**Lesson objective**: Use event-based Workflows to define the flow of information around a system\n",
    "\n",
    "In this lab, you'll learn about Workflow concepts."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42f8de14",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#fff1d7; padding:15px;\"> <b> Note</b>: Make sure to run the notebook cell by cell. Please try to avoid running all cells at once.</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a395e553-d1f1-4819-9a9b-162d6dc9cfc0",
   "metadata": {
    "height": 81
   },
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "from helper import extract_html_content\n",
    "import random\n",
    "from helper import get_openai_api_key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "80f1e17e-362a-41b5-a002-4d29b4aeca32",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "api_key = get_openai_api_key()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96e7a7f7",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#fff6ff; padding:13px; border-width:3px; border-color:#efe6ef; border-style:solid; border-radius:6px\">\n",
    "<p> 💻 &nbsp; <b>To access <code>requirements.txt</code> and <code>helper.py</code> files:</b> 1) click on the <em>\"File\"</em> option on the top menu of the notebook and then 2) click on <em>\"Open\"</em>.\n",
    "\n",
    "<p> ⬇ &nbsp; <b>Download Notebooks:</b> 1) click on the <em>\"File\"</em> option on the top menu of the notebook and then 2) click on <em>\"Download as\"</em> and select <em>\"Notebook (.ipynb)\"</em>.</p>\n",
    "\n",
    "<p> 📒 &nbsp; For more help, please see the <em>\"Appendix – Tips and Help\"</em> Lesson.</p>\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58295bb8",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#f7fff8; padding:15px; border-width:3px; border-color:#e0f0e0; border-style:solid; border-radius:6px\"> 🚨\n",
    "&nbsp; <b>Different Run Results:</b> The output generated by AI chat models can vary with each execution due to their dynamic, probabilistic nature. Don't be surprised if your results differ from those shown in the video.</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7da1d6c3-1026-4800-8add-9ad555e94310",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Creating a Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7ddba6-ce8a-4862-82b0-d28fc7719a9a",
   "metadata": {},
   "source": [
    "Under the hood, Workflows are regular Python classes. They are defined as a series of steps, each of which receives certain classes of events and emits certain classes of events.\n",
    "\n",
    "Here's the most basic form of a workflow, with a single step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7869015c-7d6e-4bc4-9f0c-34bdb221a974",
   "metadata": {
    "height": 132
   },
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Workflow,\n",
    "    step,\n",
    "    Context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "caa53423-2c92-403f-b190-b4ee646870dd",
   "metadata": {
    "height": 130
   },
   "outputs": [],
   "source": [
    "class MyWorkflow(Workflow):\n",
    "    # declare a function as a step\n",
    "    @step\n",
    "    async def my_step(self, ev: StartEvent) -> StopEvent:\n",
    "        # do something here\n",
    "        return StopEvent(result=\"Hello, world!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1f4f05d-7407-4161-b5e7-e388e727f44f",
   "metadata": {},
   "source": [
    "This new `MyWorkflow` class:\n",
    "* Uses the `@step` decorator to declare a function to be a step\n",
    "* Has a single step called `my_step` which accepts a `StartEvent`. `StartEvent` is a special event which is always generated when a workflow first runs.\n",
    "* `my_step` returns a `StopEvent`, which is another special event. When a `StopEvent` is emitted the workflow returns it and stops running.\n",
    "\n",
    "*Note*: The **async** keyword defines asynchronous functions, which can be paused and resumed, allowing other tasks to run in the meantime.\n",
    "\n",
    "You instantiate it and run your workflow like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "64df86f7-5d90-41db-a283-9b81476dbff2",
   "metadata": {
    "height": 113
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, world!\n"
     ]
    }
   ],
   "source": [
    "# instantiate the workflow\n",
    "basic_workflow = MyWorkflow(timeout=10, verbose=False)\n",
    "# run the workflow\n",
    "result = await basic_workflow.run()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c547af9-a4e7-4b10-913c-6bdacc2c2bf8",
   "metadata": {},
   "source": [
    "*Note*: \n",
    "- The **await** keyword is used with async functions to pause execution until the specific asynchronous task is complete.\n",
    "- The **timeout** argument represent the number of seconds after which the workflow execution will be halted\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f16a02fc-79c0-471b-8ccd-f25b914f3754",
   "metadata": {},
   "source": [
    "### Side Note"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95eaef1b-ac38-4255-9e39-bbded23ce17f",
   "metadata": {},
   "source": [
    "Workflows are async by default, so you use `await` to get the result of the `run` command. This will work fine in a notebook environment; in a vanilla python script you will need to import `asyncio` and wrap your code in an async function, like this:\n",
    "\n",
    "```\n",
    "async def main():\n",
    "    w = MyWorkflow(timeout=10, verbose=False)\n",
    "    result = await w.run()\n",
    "    print(result)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import asyncio\n",
    "    asyncio.run(main())\n",
    "```\n",
    "\n",
    "Since you're in a notebook right now, you won't execute the above code as it won't work!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8922ab2-a783-46d7-8c1b-73b6e370ca66",
   "metadata": {},
   "source": [
    "## Visualizing a workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec71ab6-efb7-4740-8906-bb7d2f0ad509",
   "metadata": {},
   "source": [
    "A great feature of workflows is the built-in visualizer, which you will install now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "128fc0a3-2e07-433f-96a9-1a3b123e424d",
   "metadata": {
    "height": 45
   },
   "outputs": [],
   "source": [
    "from llama_index.utils.workflow import draw_all_possible_flows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "768b444b-095f-4d73-b51d-dc0e51cc635c",
   "metadata": {
    "height": 81
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'NoneType'>\n",
      "<class 'llama_index.core.workflow.events.StopEvent'>\n",
      "workflows/basic_workflow.html\n"
     ]
    }
   ],
   "source": [
    "draw_all_possible_flows(\n",
    "    basic_workflow, \n",
    "    filename=\"workflows/basic_workflow.html\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "853fc336-6bd3-4050-b105-6c0da5df439d",
   "metadata": {},
   "source": [
    "You are provided with a helper function which displays the HTML page generated inside the notebook. On your own computer, you could just open the file in your browser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9c960b3e-f8e6-4b82-8249-191e876ce503",
   "metadata": {
    "height": 62
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       " <div style=\"width: 100%; height: 800px; overflow: hidden;\"> <html>\n",
       "    <head>\n",
       "        <meta charset=\"utf-8\">\n",
       "        \n",
       "            <script src=\"lib/bindings/utils.js\"></script>\n",
       "            <link rel=\"stylesheet\" href=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/dist/vis-network.min.css\" integrity=\"sha512-WgxfT5LWjfszlPHXRmBWHkV2eceiWTOBvrKCNbdgDYTHrT2AeLCGbF4sZlZw3UMN3WtL0tGUoIAKsu8mllg/XA==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\" />\n",
       "            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js\" integrity=\"sha512-LnvoEWDFrqGHlHmDD2101OrLcbsfkrzoSpvtSQtxK3RMnRV0eOkhhBN2dXHKRrUU8p2DGRTk35n4O8nWSVe1mQ==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\"></script>\n",
       "            \n",
       "        \n",
       "<center>\n",
       "<h1></h1>\n",
       "</center>\n",
       "\n",
       "<!-- <link rel=\"stylesheet\" href=\"../node_modules/vis/dist/vis.min.css\" type=\"text/css\" />\n",
       "<script type=\"text/javascript\" src=\"../node_modules/vis/dist/vis.js\"> </script>-->\n",
       "        <link\n",
       "          href=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/css/bootstrap.min.css\"\n",
       "          rel=\"stylesheet\"\n",
       "          integrity=\"sha384-eOJMYsd53ii+scO/bJGFsiCZc+5NDVN2yr8+0RDqr0Ql0h+rP48ckxlpbzKgwra6\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        />\n",
       "        <script\n",
       "          src=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/js/bootstrap.bundle.min.js\"\n",
       "          integrity=\"sha384-JEW9xMcG8R+pH31jmWH6WWP0WintQrMb4s7ZOdauHnUtxwoG2vI5DkLtS3qm9Ekf\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        ></script>\n",
       "\n",
       "\n",
       "        <center>\n",
       "          <h1></h1>\n",
       "        </center>\n",
       "        <style type=\"text/css\">\n",
       "\n",
       "             #mynetwork {\n",
       "                 width: 100%;\n",
       "                 height: 750px;\n",
       "                 background-color: #ffffff;\n",
       "                 border: 1px solid lightgray;\n",
       "                 position: relative;\n",
       "                 float: left;\n",
       "             }\n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "        </style>\n",
       "    </head>\n",
       "\n",
       "\n",
       "    <body>\n",
       "        <div class=\"card\" style=\"width: 100%\">\n",
       "            \n",
       "            \n",
       "            <div id=\"mynetwork\" class=\"card-body\"></div>\n",
       "        </div>\n",
       "\n",
       "        \n",
       "        \n",
       "\n",
       "        <script type=\"text/javascript\">\n",
       "\n",
       "              // initialize global variables.\n",
       "              var edges;\n",
       "              var nodes;\n",
       "              var allNodes;\n",
       "              var allEdges;\n",
       "              var nodeColors;\n",
       "              var originalNodes;\n",
       "              var network;\n",
       "              var container;\n",
       "              var options, data;\n",
       "              var filter = {\n",
       "                  item : '',\n",
       "                  property : '',\n",
       "                  value : []\n",
       "              };\n",
       "\n",
       "              \n",
       "\n",
       "              \n",
       "\n",
       "              // This method is responsible for drawing the graph, returns the drawn network\n",
       "              function drawGraph() {\n",
       "                  var container = document.getElementById('mynetwork');\n",
       "\n",
       "                  \n",
       "\n",
       "                  // parsing and collecting nodes and edges from the python\n",
       "                  nodes = new vis.DataSet([{\"color\": \"#FFA07A\", \"id\": \"StopEvent\", \"label\": \"StopEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"_done\", \"label\": \"_done\", \"shape\": \"box\"}, {\"color\": \"#ADD8E6\", \"id\": \"my_step\", \"label\": \"my_step\", \"shape\": \"box\"}, {\"color\": \"#E27AFF\", \"id\": \"StartEvent\", \"label\": \"StartEvent\", \"shape\": \"ellipse\"}]);\n",
       "                  edges = new vis.DataSet([{\"arrows\": \"to\", \"from\": \"StopEvent\", \"to\": \"_done\"}, {\"arrows\": \"to\", \"from\": \"my_step\", \"to\": \"StopEvent\"}, {\"arrows\": \"to\", \"from\": \"StartEvent\", \"to\": \"my_step\"}]);\n",
       "\n",
       "                  nodeColors = {};\n",
       "                  allNodes = nodes.get({ returnType: \"Object\" });\n",
       "                  for (nodeId in allNodes) {\n",
       "                    nodeColors[nodeId] = allNodes[nodeId].color;\n",
       "                  }\n",
       "                  allEdges = edges.get({ returnType: \"Object\" });\n",
       "                  // adding nodes and edges to the graph\n",
       "                  data = {nodes: nodes, edges: edges};\n",
       "\n",
       "                  var options = {\n",
       "    \"configure\": {\n",
       "        \"enabled\": false\n",
       "    },\n",
       "    \"edges\": {\n",
       "        \"color\": {\n",
       "            \"inherit\": true\n",
       "        },\n",
       "        \"smooth\": {\n",
       "            \"enabled\": true,\n",
       "            \"type\": \"dynamic\"\n",
       "        }\n",
       "    },\n",
       "    \"interaction\": {\n",
       "        \"dragNodes\": true,\n",
       "        \"hideEdgesOnDrag\": false,\n",
       "        \"hideNodesOnDrag\": false\n",
       "    },\n",
       "    \"physics\": {\n",
       "        \"enabled\": true,\n",
       "        \"stabilization\": {\n",
       "            \"enabled\": true,\n",
       "            \"fit\": true,\n",
       "            \"iterations\": 1000,\n",
       "            \"onlyDynamicEdges\": false,\n",
       "            \"updateInterval\": 50\n",
       "        }\n",
       "    }\n",
       "};\n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  network = new vis.Network(container, data, options);\n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  return network;\n",
       "\n",
       "              }\n",
       "              drawGraph();\n",
       "        </script>\n",
       "    </body>\n",
       "</html> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "isolated": true
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "html_content = extract_html_content(\"workflows/basic_workflow.html\")\n",
    "display(HTML(html_content), metadata=dict(isolated=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b07b66b3",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#fff1d7; padding:15px;\"> <b> Note </b>: The visualized workflow might look slightly different from that of the video. If it doesn't display in the notebook, you can: 1) click on the <em>\"File\"</em> option on the top menu of the notebook and then 2) click on <em>\"Open\"</em>. The HTML file is in the folder \"workflows\".</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54381b5b-5007-422b-bbc8-635628add88d",
   "metadata": {},
   "source": [
    "Of course, a flow with a single step is not very useful! Let's define a multi-step workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23f2fe8c-26f8-465b-b099-3744982bb5f8",
   "metadata": {},
   "source": [
    "## Creating Custom Events"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d19c242a-a249-437c-86a2-9de7f43b4344",
   "metadata": {},
   "source": [
    "Multiple steps are created by defining custom events that can be emitted by steps and trigger other steps. Let's define a simple 3-step workflow by defining two custom events, `FirstEvent` and `SecondEvent`. These classes can have any names and properties, but must inherit from `Event`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f42526f9-8726-4b8a-a99e-860c3fb911ed",
   "metadata": {
    "height": 132
   },
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Event\n",
    "\n",
    "class FirstEvent(Event):\n",
    "    first_output: str\n",
    "\n",
    "class SecondEvent(Event):\n",
    "    second_output: str"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f9a8356-379c-4cb7-ad6e-99543fe0f1c0",
   "metadata": {},
   "source": [
    "### Defining the workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25a7ba9c-2119-463f-bbf5-1207349e9429",
   "metadata": {},
   "source": [
    "Now you define the workflow itself. You do this by defining the input and output types on each step.\n",
    "\n",
    "* `step_one` takes a `StartEvent` and returns a `FirstEvent`\n",
    "* `step_two` takes a `FirstEvent` and returns a `SecondEvent`\n",
    "* `step_three` takes a `SecondEvent` and returns a `StopEvent`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1fb67470-dd94-4a91-b040-e243928a788e",
   "metadata": {
    "height": 283
   },
   "outputs": [],
   "source": [
    "class MyWorkflow(Workflow):\n",
    "    @step\n",
    "    async def step_one(self, ev: StartEvent) -> FirstEvent:\n",
    "        print(ev.first_input)\n",
    "        return FirstEvent(first_output=\"First step complete.\")\n",
    "\n",
    "    @step\n",
    "    async def step_two(self, ev: FirstEvent) -> SecondEvent:\n",
    "        print(ev.first_output)\n",
    "        return SecondEvent(second_output=\"Second step complete.\")\n",
    "\n",
    "    @step\n",
    "    async def step_three(self, ev: SecondEvent) -> StopEvent:\n",
    "        print(ev.second_output)\n",
    "        return StopEvent(result=\"Workflow complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3b55bd-f41c-44bc-b54f-93cdf51b790c",
   "metadata": {},
   "source": [
    "*Note:* Properties of `StartEvent` and `StopEvent`:\n",
    "- For `StartEvent`, you define its properties and pass in their values when you run the workflow as shown in the next cell.\n",
    "- For `StopEvent`, by default, it only has one property `result`. You can always create a class that inherits `StopEvent` so you can customize what it returns."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f61269cd-0dea-4cb0-9ef8-a248b42dcb8b",
   "metadata": {},
   "source": [
    "You run this just like you ran the other workflows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "85399e8e-269c-42e4-b0be-1e7f9bdcf67e",
   "metadata": {
    "height": 79
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start the workflow.\n",
      "First step complete.\n",
      "Second step complete.\n",
      "Workflow complete.\n"
     ]
    }
   ],
   "source": [
    "workflow = MyWorkflow(timeout=10, verbose=False)\n",
    "result = await workflow.run(first_input=\"Start the workflow.\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b02fcb-acb7-40d6-a137-3857bb16a77b",
   "metadata": {},
   "source": [
    "And you can visualize it just like you did before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e88ff080-570d-43f3-b189-07108f6902d5",
   "metadata": {
    "height": 62
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'NoneType'>\n",
      "<class '__main__.FirstEvent'>\n",
      "<class 'llama_index.core.workflow.events.StopEvent'>\n",
      "<class '__main__.SecondEvent'>\n",
      "workflows/custom_events.html\n"
     ]
    }
   ],
   "source": [
    "WORKFLOW_FILE = \"workflows/custom_events.html\"\n",
    "draw_all_possible_flows(workflow, filename=WORKFLOW_FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "aaa85bd8-7d70-428e-a19d-a3e1a4ca776d",
   "metadata": {
    "height": 62
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       " <div style=\"width: 100%; height: 800px; overflow: hidden;\"> <html>\n",
       "    <head>\n",
       "        <meta charset=\"utf-8\">\n",
       "        \n",
       "            <script src=\"lib/bindings/utils.js\"></script>\n",
       "            <link rel=\"stylesheet\" href=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/dist/vis-network.min.css\" integrity=\"sha512-WgxfT5LWjfszlPHXRmBWHkV2eceiWTOBvrKCNbdgDYTHrT2AeLCGbF4sZlZw3UMN3WtL0tGUoIAKsu8mllg/XA==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\" />\n",
       "            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js\" integrity=\"sha512-LnvoEWDFrqGHlHmDD2101OrLcbsfkrzoSpvtSQtxK3RMnRV0eOkhhBN2dXHKRrUU8p2DGRTk35n4O8nWSVe1mQ==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\"></script>\n",
       "            \n",
       "        \n",
       "<center>\n",
       "<h1></h1>\n",
       "</center>\n",
       "\n",
       "<!-- <link rel=\"stylesheet\" href=\"../node_modules/vis/dist/vis.min.css\" type=\"text/css\" />\n",
       "<script type=\"text/javascript\" src=\"../node_modules/vis/dist/vis.js\"> </script>-->\n",
       "        <link\n",
       "          href=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/css/bootstrap.min.css\"\n",
       "          rel=\"stylesheet\"\n",
       "          integrity=\"sha384-eOJMYsd53ii+scO/bJGFsiCZc+5NDVN2yr8+0RDqr0Ql0h+rP48ckxlpbzKgwra6\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        />\n",
       "        <script\n",
       "          src=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/js/bootstrap.bundle.min.js\"\n",
       "          integrity=\"sha384-JEW9xMcG8R+pH31jmWH6WWP0WintQrMb4s7ZOdauHnUtxwoG2vI5DkLtS3qm9Ekf\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        ></script>\n",
       "\n",
       "\n",
       "        <center>\n",
       "          <h1></h1>\n",
       "        </center>\n",
       "        <style type=\"text/css\">\n",
       "\n",
       "             #mynetwork {\n",
       "                 width: 100%;\n",
       "                 height: 750px;\n",
       "                 background-color: #ffffff;\n",
       "                 border: 1px solid lightgray;\n",
       "                 position: relative;\n",
       "                 float: left;\n",
       "             }\n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "        </style>\n",
       "    </head>\n",
       "\n",
       "\n",
       "    <body>\n",
       "        <div class=\"card\" style=\"width: 100%\">\n",
       "            \n",
       "            \n",
       "            <div id=\"mynetwork\" class=\"card-body\"></div>\n",
       "        </div>\n",
       "\n",
       "        \n",
       "        \n",
       "\n",
       "        <script type=\"text/javascript\">\n",
       "\n",
       "              // initialize global variables.\n",
       "              var edges;\n",
       "              var nodes;\n",
       "              var allNodes;\n",
       "              var allEdges;\n",
       "              var nodeColors;\n",
       "              var originalNodes;\n",
       "              var network;\n",
       "              var container;\n",
       "              var options, data;\n",
       "              var filter = {\n",
       "                  item : '',\n",
       "                  property : '',\n",
       "                  value : []\n",
       "              };\n",
       "\n",
       "              \n",
       "\n",
       "              \n",
       "\n",
       "              // This method is responsible for drawing the graph, returns the drawn network\n",
       "              function drawGraph() {\n",
       "                  var container = document.getElementById('mynetwork');\n",
       "\n",
       "                  \n",
       "\n",
       "                  // parsing and collecting nodes and edges from the python\n",
       "                  nodes = new vis.DataSet([{\"color\": \"#FFA07A\", \"id\": \"StopEvent\", \"label\": \"StopEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"_done\", \"label\": \"_done\", \"shape\": \"box\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_one\", \"label\": \"step_one\", \"shape\": \"box\"}, {\"color\": \"#E27AFF\", \"id\": \"StartEvent\", \"label\": \"StartEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_three\", \"label\": \"step_three\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"SecondEvent\", \"label\": \"SecondEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_two\", \"label\": \"step_two\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"FirstEvent\", \"label\": \"FirstEvent\", \"shape\": \"ellipse\"}]);\n",
       "                  edges = new vis.DataSet([{\"arrows\": \"to\", \"from\": \"StopEvent\", \"to\": \"_done\"}, {\"arrows\": \"to\", \"from\": \"step_one\", \"to\": \"FirstEvent\"}, {\"arrows\": \"to\", \"from\": \"StartEvent\", \"to\": \"step_one\"}, {\"arrows\": \"to\", \"from\": \"step_three\", \"to\": \"StopEvent\"}, {\"arrows\": \"to\", \"from\": \"SecondEvent\", \"to\": \"step_three\"}, {\"arrows\": \"to\", \"from\": \"step_two\", \"to\": \"SecondEvent\"}, {\"arrows\": \"to\", \"from\": \"FirstEvent\", \"to\": \"step_two\"}]);\n",
       "\n",
       "                  nodeColors = {};\n",
       "                  allNodes = nodes.get({ returnType: \"Object\" });\n",
       "                  for (nodeId in allNodes) {\n",
       "                    nodeColors[nodeId] = allNodes[nodeId].color;\n",
       "                  }\n",
       "                  allEdges = edges.get({ returnType: \"Object\" });\n",
       "                  // adding nodes and edges to the graph\n",
       "                  data = {nodes: nodes, edges: edges};\n",
       "\n",
       "                  var options = {\n",
       "    \"configure\": {\n",
       "        \"enabled\": false\n",
       "    },\n",
       "    \"edges\": {\n",
       "        \"color\": {\n",
       "            \"inherit\": true\n",
       "        },\n",
       "        \"smooth\": {\n",
       "            \"enabled\": true,\n",
       "            \"type\": \"dynamic\"\n",
       "        }\n",
       "    },\n",
       "    \"interaction\": {\n",
       "        \"dragNodes\": true,\n",
       "        \"hideEdgesOnDrag\": false,\n",
       "        \"hideNodesOnDrag\": false\n",
       "    },\n",
       "    \"physics\": {\n",
       "        \"enabled\": true,\n",
       "        \"stabilization\": {\n",
       "            \"enabled\": true,\n",
       "            \"fit\": true,\n",
       "            \"iterations\": 1000,\n",
       "            \"onlyDynamicEdges\": false,\n",
       "            \"updateInterval\": 50\n",
       "        }\n",
       "    }\n",
       "};\n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  network = new vis.Network(container, data, options);\n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  return network;\n",
       "\n",
       "              }\n",
       "              drawGraph();\n",
       "        </script>\n",
       "    </body>\n",
       "</html> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "isolated": true
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "html_content = extract_html_content(WORKFLOW_FILE)\n",
    "display(HTML(html_content), metadata=dict(isolated=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "610befa4-d934-4035-9a9f-68593b716b7a",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Creating Loops"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0651e4e-6aeb-47a5-a028-47c8910521bd",
   "metadata": {},
   "source": [
    "However, there's not much point to a workflow if it just runs straight through! A key feature of Workflows is their enablement of branching and looping logic, more simply and flexibly than graph-based approaches. To enable looping, let's create a new `LoopEvent`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4b2d2c6f-a2f6-4321-986e-7c7c0a01af75",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "class LoopEvent(Event):\n",
    "    loop_output: str"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30efc66f-d647-4582-9db9-c66f2a9e80ed",
   "metadata": {},
   "source": [
    "Now you'll edit your `step_one` to make a random decision about whether to execute serially or loop back:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "28458e79-7022-4e66-934d-53ea2845283f",
   "metadata": {
    "height": 351
   },
   "outputs": [],
   "source": [
    "class MyWorkflow(Workflow):\n",
    "    @step\n",
    "    async def step_one(self, ev: StartEvent | LoopEvent) -> FirstEvent | LoopEvent:\n",
    "        if random.randint(0, 1) == 0:\n",
    "            print(\"Bad thing happened\")\n",
    "            return LoopEvent(loop_output=\"Back to step one.\")\n",
    "        else:\n",
    "            print(\"Good thing happened\")\n",
    "            return FirstEvent(first_output=\"First step complete.\")\n",
    "\n",
    "    @step\n",
    "    async def step_two(self, ev: FirstEvent) -> SecondEvent:\n",
    "        print(ev.first_output)\n",
    "        return SecondEvent(second_output=\"Second step complete.\")\n",
    "\n",
    "    @step\n",
    "    async def step_three(self, ev: SecondEvent) -> StopEvent:\n",
    "        print(ev.second_output)\n",
    "        return StopEvent(result=\"Workflow complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d5ff1e1-1646-4e30-a8e1-36b5bfe245f9",
   "metadata": {},
   "source": [
    "Note the new type annotations on `step_one`: the step now accepts either a `StartEvent` or a `LoopEvent` to trigger the step, and it also emits either a `FirstEvent` or a `LoopEvent`.\n",
    "\n",
    "You run it as usual. You might need to run it a couple of times to see the loop happen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a9a27c0f-9661-4f06-a8e2-57308056c115",
   "metadata": {
    "height": 79
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bad thing happened\n",
      "Bad thing happened\n",
      "Bad thing happened\n",
      "Good thing happened\n",
      "First step complete.\n",
      "Second step complete.\n",
      "Workflow complete.\n"
     ]
    }
   ],
   "source": [
    "loop_workflow = MyWorkflow(timeout=10, verbose=False)\n",
    "result = await loop_workflow.run(first_input=\"Start the workflow.\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c9164a7-767f-4815-bd79-80d60b0eff0a",
   "metadata": {},
   "source": [
    "Your new, looping workflow visualizes like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "357849c3-f0eb-4285-a4b7-0b1c5c5205fc",
   "metadata": {
    "height": 96
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'NoneType'>\n",
      "<class '__main__.FirstEvent'>\n",
      "<class '__main__.LoopEvent'>\n",
      "<class 'llama_index.core.workflow.events.StopEvent'>\n",
      "<class '__main__.SecondEvent'>\n",
      "workflows/loop_events.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       " <div style=\"width: 100%; height: 800px; overflow: hidden;\"> <html>\n",
       "    <head>\n",
       "        <meta charset=\"utf-8\">\n",
       "        \n",
       "            <script src=\"lib/bindings/utils.js\"></script>\n",
       "            <link rel=\"stylesheet\" href=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/dist/vis-network.min.css\" integrity=\"sha512-WgxfT5LWjfszlPHXRmBWHkV2eceiWTOBvrKCNbdgDYTHrT2AeLCGbF4sZlZw3UMN3WtL0tGUoIAKsu8mllg/XA==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\" />\n",
       "            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js\" integrity=\"sha512-LnvoEWDFrqGHlHmDD2101OrLcbsfkrzoSpvtSQtxK3RMnRV0eOkhhBN2dXHKRrUU8p2DGRTk35n4O8nWSVe1mQ==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\"></script>\n",
       "            \n",
       "        \n",
       "<center>\n",
       "<h1></h1>\n",
       "</center>\n",
       "\n",
       "<!-- <link rel=\"stylesheet\" href=\"../node_modules/vis/dist/vis.min.css\" type=\"text/css\" />\n",
       "<script type=\"text/javascript\" src=\"../node_modules/vis/dist/vis.js\"> </script>-->\n",
       "        <link\n",
       "          href=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/css/bootstrap.min.css\"\n",
       "          rel=\"stylesheet\"\n",
       "          integrity=\"sha384-eOJMYsd53ii+scO/bJGFsiCZc+5NDVN2yr8+0RDqr0Ql0h+rP48ckxlpbzKgwra6\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        />\n",
       "        <script\n",
       "          src=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/js/bootstrap.bundle.min.js\"\n",
       "          integrity=\"sha384-JEW9xMcG8R+pH31jmWH6WWP0WintQrMb4s7ZOdauHnUtxwoG2vI5DkLtS3qm9Ekf\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        ></script>\n",
       "\n",
       "\n",
       "        <center>\n",
       "          <h1></h1>\n",
       "        </center>\n",
       "        <style type=\"text/css\">\n",
       "\n",
       "             #mynetwork {\n",
       "                 width: 100%;\n",
       "                 height: 750px;\n",
       "                 background-color: #ffffff;\n",
       "                 border: 1px solid lightgray;\n",
       "                 position: relative;\n",
       "                 float: left;\n",
       "             }\n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "        </style>\n",
       "    </head>\n",
       "\n",
       "\n",
       "    <body>\n",
       "        <div class=\"card\" style=\"width: 100%\">\n",
       "            \n",
       "            \n",
       "            <div id=\"mynetwork\" class=\"card-body\"></div>\n",
       "        </div>\n",
       "\n",
       "        \n",
       "        \n",
       "\n",
       "        <script type=\"text/javascript\">\n",
       "\n",
       "              // initialize global variables.\n",
       "              var edges;\n",
       "              var nodes;\n",
       "              var allNodes;\n",
       "              var allEdges;\n",
       "              var nodeColors;\n",
       "              var originalNodes;\n",
       "              var network;\n",
       "              var container;\n",
       "              var options, data;\n",
       "              var filter = {\n",
       "                  item : '',\n",
       "                  property : '',\n",
       "                  value : []\n",
       "              };\n",
       "\n",
       "              \n",
       "\n",
       "              \n",
       "\n",
       "              // This method is responsible for drawing the graph, returns the drawn network\n",
       "              function drawGraph() {\n",
       "                  var container = document.getElementById('mynetwork');\n",
       "\n",
       "                  \n",
       "\n",
       "                  // parsing and collecting nodes and edges from the python\n",
       "                  nodes = new vis.DataSet([{\"color\": \"#FFA07A\", \"id\": \"StopEvent\", \"label\": \"StopEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"_done\", \"label\": \"_done\", \"shape\": \"box\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_one\", \"label\": \"step_one\", \"shape\": \"box\"}, {\"color\": \"#E27AFF\", \"id\": \"StartEvent\", \"label\": \"StartEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#90EE90\", \"id\": \"LoopEvent\", \"label\": \"LoopEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_three\", \"label\": \"step_three\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"SecondEvent\", \"label\": \"SecondEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_two\", \"label\": \"step_two\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"FirstEvent\", \"label\": \"FirstEvent\", \"shape\": \"ellipse\"}]);\n",
       "                  edges = new vis.DataSet([{\"arrows\": \"to\", \"from\": \"StopEvent\", \"to\": \"_done\"}, {\"arrows\": \"to\", \"from\": \"step_one\", \"to\": \"FirstEvent\"}, {\"arrows\": \"to\", \"from\": \"step_one\", \"to\": \"LoopEvent\"}, {\"arrows\": \"to\", \"from\": \"StartEvent\", \"to\": \"step_one\"}, {\"arrows\": \"to\", \"from\": \"LoopEvent\", \"to\": \"step_one\"}, {\"arrows\": \"to\", \"from\": \"step_three\", \"to\": \"StopEvent\"}, {\"arrows\": \"to\", \"from\": \"SecondEvent\", \"to\": \"step_three\"}, {\"arrows\": \"to\", \"from\": \"step_two\", \"to\": \"SecondEvent\"}, {\"arrows\": \"to\", \"from\": \"FirstEvent\", \"to\": \"step_two\"}]);\n",
       "\n",
       "                  nodeColors = {};\n",
       "                  allNodes = nodes.get({ returnType: \"Object\" });\n",
       "                  for (nodeId in allNodes) {\n",
       "                    nodeColors[nodeId] = allNodes[nodeId].color;\n",
       "                  }\n",
       "                  allEdges = edges.get({ returnType: \"Object\" });\n",
       "                  // adding nodes and edges to the graph\n",
       "                  data = {nodes: nodes, edges: edges};\n",
       "\n",
       "                  var options = {\n",
       "    \"configure\": {\n",
       "        \"enabled\": false\n",
       "    },\n",
       "    \"edges\": {\n",
       "        \"color\": {\n",
       "            \"inherit\": true\n",
       "        },\n",
       "        \"smooth\": {\n",
       "            \"enabled\": true,\n",
       "            \"type\": \"dynamic\"\n",
       "        }\n",
       "    },\n",
       "    \"interaction\": {\n",
       "        \"dragNodes\": true,\n",
       "        \"hideEdgesOnDrag\": false,\n",
       "        \"hideNodesOnDrag\": false\n",
       "    },\n",
       "    \"physics\": {\n",
       "        \"enabled\": true,\n",
       "        \"stabilization\": {\n",
       "            \"enabled\": true,\n",
       "            \"fit\": true,\n",
       "            \"iterations\": 1000,\n",
       "            \"onlyDynamicEdges\": false,\n",
       "            \"updateInterval\": 50\n",
       "        }\n",
       "    }\n",
       "};\n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  network = new vis.Network(container, data, options);\n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  return network;\n",
       "\n",
       "              }\n",
       "              drawGraph();\n",
       "        </script>\n",
       "    </body>\n",
       "</html> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "isolated": true
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "WORKFLOW_FILE = \"workflows/loop_events.html\"\n",
    "draw_all_possible_flows(loop_workflow, filename=WORKFLOW_FILE)\n",
    "html_content = extract_html_content(WORKFLOW_FILE)\n",
    "display(HTML(html_content), metadata=dict(isolated=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdebabcb-e536-4fa0-8354-df1bac208543",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Branching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b23d594e-c4e4-42c6-a88f-2f8ab7885826",
   "metadata": {},
   "source": [
    "The same constructs that allow you to loop allow you to create branches. Here's a workflow that executes two different branches depending on an early decision:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6d378818-6dc1-4994-9729-bc34399b1147",
   "metadata": {
    "height": 200
   },
   "outputs": [],
   "source": [
    "class BranchA1Event(Event):\n",
    "    payload: str\n",
    "\n",
    "class BranchA2Event(Event):\n",
    "    payload: str\n",
    "\n",
    "class BranchB1Event(Event):\n",
    "    payload: str\n",
    "\n",
    "class BranchB2Event(Event):\n",
    "    payload: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ce851bf7-6125-4fbb-9d06-4e30fc3283b1",
   "metadata": {
    "height": 521
   },
   "outputs": [],
   "source": [
    "class BranchWorkflow(Workflow):\n",
    "    @step\n",
    "    async def start(self, ev: StartEvent) -> BranchA1Event | BranchB1Event:\n",
    "        if random.randint(0, 1) == 0:\n",
    "            print(\"Go to branch A\")\n",
    "            return BranchA1Event(payload=\"Branch A\")\n",
    "        else:\n",
    "            print(\"Go to branch B\")\n",
    "            return BranchB1Event(payload=\"Branch B\")\n",
    "\n",
    "    @step\n",
    "    async def step_a1(self, ev: BranchA1Event) -> BranchA2Event:\n",
    "        print(ev.payload)\n",
    "        return BranchA2Event(payload=ev.payload)\n",
    "\n",
    "    @step\n",
    "    async def step_b1(self, ev: BranchB1Event) -> BranchB2Event:\n",
    "        print(ev.payload)\n",
    "        return BranchB2Event(payload=ev.payload)\n",
    "\n",
    "    @step\n",
    "    async def step_a2(self, ev: BranchA2Event) -> StopEvent:\n",
    "        print(ev.payload)\n",
    "        return StopEvent(result=\"Branch A complete.\")\n",
    "\n",
    "    @step\n",
    "    async def step_b2(self, ev: BranchB2Event) -> StopEvent:\n",
    "        print(ev.payload)\n",
    "        return StopEvent(result=\"Branch B complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4919c59-75c4-4fa8-a10d-d63f5a3fd1f3",
   "metadata": {},
   "source": [
    "You don't actually need to instantiate the workflow to visualize it, you can just pass the workflow class directly to the visualizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "911c2e4f-3ead-4811-8d63-efd706f2d020",
   "metadata": {
    "height": 79
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'NoneType'>\n",
      "<class '__main__.BranchA1Event'>\n",
      "<class '__main__.BranchB1Event'>\n",
      "<class '__main__.BranchA2Event'>\n",
      "<class 'llama_index.core.workflow.events.StopEvent'>\n",
      "<class '__main__.BranchB2Event'>\n",
      "<class 'llama_index.core.workflow.events.StopEvent'>\n",
      "workflows/branching.html\n"
     ]
    }
   ],
   "source": [
    "WORKFLOW_FILE = \"workflows/branching.html\"\n",
    "draw_all_possible_flows(BranchWorkflow, filename=WORKFLOW_FILE)\n",
    "html_content = extract_html_content(WORKFLOW_FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e7a64ca9-50b3-4e22-a96f-5f3bfe57bf79",
   "metadata": {
    "height": 45
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       " <div style=\"width: 100%; height: 800px; overflow: hidden;\"> <html>\n",
       "    <head>\n",
       "        <meta charset=\"utf-8\">\n",
       "        \n",
       "            <script src=\"lib/bindings/utils.js\"></script>\n",
       "            <link rel=\"stylesheet\" href=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/dist/vis-network.min.css\" integrity=\"sha512-WgxfT5LWjfszlPHXRmBWHkV2eceiWTOBvrKCNbdgDYTHrT2AeLCGbF4sZlZw3UMN3WtL0tGUoIAKsu8mllg/XA==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\" />\n",
       "            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js\" integrity=\"sha512-LnvoEWDFrqGHlHmDD2101OrLcbsfkrzoSpvtSQtxK3RMnRV0eOkhhBN2dXHKRrUU8p2DGRTk35n4O8nWSVe1mQ==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\"></script>\n",
       "            \n",
       "        \n",
       "<center>\n",
       "<h1></h1>\n",
       "</center>\n",
       "\n",
       "<!-- <link rel=\"stylesheet\" href=\"../node_modules/vis/dist/vis.min.css\" type=\"text/css\" />\n",
       "<script type=\"text/javascript\" src=\"../node_modules/vis/dist/vis.js\"> </script>-->\n",
       "        <link\n",
       "          href=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/css/bootstrap.min.css\"\n",
       "          rel=\"stylesheet\"\n",
       "          integrity=\"sha384-eOJMYsd53ii+scO/bJGFsiCZc+5NDVN2yr8+0RDqr0Ql0h+rP48ckxlpbzKgwra6\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        />\n",
       "        <script\n",
       "          src=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/js/bootstrap.bundle.min.js\"\n",
       "          integrity=\"sha384-JEW9xMcG8R+pH31jmWH6WWP0WintQrMb4s7ZOdauHnUtxwoG2vI5DkLtS3qm9Ekf\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        ></script>\n",
       "\n",
       "\n",
       "        <center>\n",
       "          <h1></h1>\n",
       "        </center>\n",
       "        <style type=\"text/css\">\n",
       "\n",
       "             #mynetwork {\n",
       "                 width: 100%;\n",
       "                 height: 750px;\n",
       "                 background-color: #ffffff;\n",
       "                 border: 1px solid lightgray;\n",
       "                 position: relative;\n",
       "                 float: left;\n",
       "             }\n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "        </style>\n",
       "    </head>\n",
       "\n",
       "\n",
       "    <body>\n",
       "        <div class=\"card\" style=\"width: 100%\">\n",
       "            \n",
       "            \n",
       "            <div id=\"mynetwork\" class=\"card-body\"></div>\n",
       "        </div>\n",
       "\n",
       "        \n",
       "        \n",
       "\n",
       "        <script type=\"text/javascript\">\n",
       "\n",
       "              // initialize global variables.\n",
       "              var edges;\n",
       "              var nodes;\n",
       "              var allNodes;\n",
       "              var allEdges;\n",
       "              var nodeColors;\n",
       "              var originalNodes;\n",
       "              var network;\n",
       "              var container;\n",
       "              var options, data;\n",
       "              var filter = {\n",
       "                  item : '',\n",
       "                  property : '',\n",
       "                  value : []\n",
       "              };\n",
       "\n",
       "              \n",
       "\n",
       "              \n",
       "\n",
       "              // This method is responsible for drawing the graph, returns the drawn network\n",
       "              function drawGraph() {\n",
       "                  var container = document.getElementById('mynetwork');\n",
       "\n",
       "                  \n",
       "\n",
       "                  // parsing and collecting nodes and edges from the python\n",
       "                  nodes = new vis.DataSet([{\"color\": \"#FFA07A\", \"id\": \"StopEvent\", \"label\": \"StopEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"_done\", \"label\": \"_done\", \"shape\": \"box\"}, {\"color\": \"#ADD8E6\", \"id\": \"start\", \"label\": \"start\", \"shape\": \"box\"}, {\"color\": \"#E27AFF\", \"id\": \"StartEvent\", \"label\": \"StartEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_a1\", \"label\": \"step_a1\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"BranchA1Event\", \"label\": \"BranchA1Event\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_a2\", \"label\": \"step_a2\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"BranchA2Event\", \"label\": \"BranchA2Event\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_b1\", \"label\": \"step_b1\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"BranchB1Event\", \"label\": \"BranchB1Event\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_b2\", \"label\": \"step_b2\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"BranchB2Event\", \"label\": \"BranchB2Event\", \"shape\": \"ellipse\"}]);\n",
       "                  edges = new vis.DataSet([{\"arrows\": \"to\", \"from\": \"StopEvent\", \"to\": \"_done\"}, {\"arrows\": \"to\", \"from\": \"start\", \"to\": \"BranchA1Event\"}, {\"arrows\": \"to\", \"from\": \"start\", \"to\": \"BranchB1Event\"}, {\"arrows\": \"to\", \"from\": \"StartEvent\", \"to\": \"start\"}, {\"arrows\": \"to\", \"from\": \"step_a1\", \"to\": \"BranchA2Event\"}, {\"arrows\": \"to\", \"from\": \"BranchA1Event\", \"to\": \"step_a1\"}, {\"arrows\": \"to\", \"from\": \"step_a2\", \"to\": \"StopEvent\"}, {\"arrows\": \"to\", \"from\": \"BranchA2Event\", \"to\": \"step_a2\"}, {\"arrows\": \"to\", \"from\": \"step_b1\", \"to\": \"BranchB2Event\"}, {\"arrows\": \"to\", \"from\": \"BranchB1Event\", \"to\": \"step_b1\"}, {\"arrows\": \"to\", \"from\": \"step_b2\", \"to\": \"StopEvent\"}, {\"arrows\": \"to\", \"from\": \"BranchB2Event\", \"to\": \"step_b2\"}]);\n",
       "\n",
       "                  nodeColors = {};\n",
       "                  allNodes = nodes.get({ returnType: \"Object\" });\n",
       "                  for (nodeId in allNodes) {\n",
       "                    nodeColors[nodeId] = allNodes[nodeId].color;\n",
       "                  }\n",
       "                  allEdges = edges.get({ returnType: \"Object\" });\n",
       "                  // adding nodes and edges to the graph\n",
       "                  data = {nodes: nodes, edges: edges};\n",
       "\n",
       "                  var options = {\n",
       "    \"configure\": {\n",
       "        \"enabled\": false\n",
       "    },\n",
       "    \"edges\": {\n",
       "        \"color\": {\n",
       "            \"inherit\": true\n",
       "        },\n",
       "        \"smooth\": {\n",
       "            \"enabled\": true,\n",
       "            \"type\": \"dynamic\"\n",
       "        }\n",
       "    },\n",
       "    \"interaction\": {\n",
       "        \"dragNodes\": true,\n",
       "        \"hideEdgesOnDrag\": false,\n",
       "        \"hideNodesOnDrag\": false\n",
       "    },\n",
       "    \"physics\": {\n",
       "        \"enabled\": true,\n",
       "        \"stabilization\": {\n",
       "            \"enabled\": true,\n",
       "            \"fit\": true,\n",
       "            \"iterations\": 1000,\n",
       "            \"onlyDynamicEdges\": false,\n",
       "            \"updateInterval\": 50\n",
       "        }\n",
       "    }\n",
       "};\n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  network = new vis.Network(container, data, options);\n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  return network;\n",
       "\n",
       "              }\n",
       "              drawGraph();\n",
       "        </script>\n",
       "    </body>\n",
       "</html> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "isolated": true
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(html_content), metadata=dict(isolated=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45090b9a-eb7a-4cf6-971a-4014621e5a51",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Concurent Execution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cdd4534-897f-497d-9123-172598073dd0",
   "metadata": {},
   "source": [
    "The final form of flow control you can implement in workflows is concurrent execution. This allows you to efficiently run long-running tasks in parallel, and gather them together when they are needed. Let's see how this is done.\n",
    "\n",
    "You'll be using a new concept, the `Context` object. This is a form of shared memory available to every step in a workflow: to access it, declare it as an argument to your step and it will be automatically populated.\n",
    "\n",
    "In this example, you'll use `Context.send_event` rather than returning an event. This allows you to emit multiple events in parallel rather than returning just one as you did previously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1b487882-49a0-42ae-b862-febde9e67e01",
   "metadata": {
    "height": 334
   },
   "outputs": [],
   "source": [
    "import asyncio\n",
    "\n",
    "class StepTwoEvent(Event):\n",
    "    query: str\n",
    "\n",
    "class ParallelFlow(Workflow):\n",
    "    @step\n",
    "    async def start(self, ctx: Context, ev: StartEvent) -> StepTwoEvent:\n",
    "        ctx.send_event(StepTwoEvent(query=\"Query 1\"))\n",
    "        ctx.send_event(StepTwoEvent(query=\"Query 2\"))\n",
    "        ctx.send_event(StepTwoEvent(query=\"Query 3\"))\n",
    "\n",
    "    @step(num_workers=4)\n",
    "    async def step_two(self, ctx: Context, ev: StepTwoEvent) -> StopEvent:\n",
    "        print(\"Running slow query \", ev.query)\n",
    "        await asyncio.sleep(random.randint(1, 5))\n",
    "\n",
    "        return StopEvent(result=ev.query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "8ec6f9a6-6585-4c7a-9268-a47b52550d37",
   "metadata": {
    "height": 79
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running slow query  Query 1\n",
      "Running slow query  Query 2\n",
      "Running slow query  Query 3\n",
      "Query 1\n"
     ]
    }
   ],
   "source": [
    "parallel_workflow = ParallelFlow(timeout=10, verbose=False)\n",
    "result = await parallel_workflow.run()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f4f00ab-639f-41ac-896b-59705e57f5d7",
   "metadata": {},
   "source": [
    "### Collecting events"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a992084-884b-4a58-bcbd-a7e4b2e14e3d",
   "metadata": {},
   "source": [
    "But what if you do want the output of all 3 events? Another method, `Context.collect_events`, exists for that purpose:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ab2b25fb-5d7d-4a0e-b270-93f03a1e11ab",
   "metadata": {
    "height": 487
   },
   "outputs": [],
   "source": [
    "class StepThreeEvent(Event):\n",
    "    result: str\n",
    "\n",
    "class ConcurrentFlow(Workflow):\n",
    "    @step\n",
    "    async def start(self, ctx: Context, ev: StartEvent) -> StepTwoEvent:\n",
    "        ctx.send_event(StepTwoEvent(query=\"Query 1\"))\n",
    "        ctx.send_event(StepTwoEvent(query=\"Query 2\"))\n",
    "        ctx.send_event(StepTwoEvent(query=\"Query 3\"))\n",
    "\n",
    "    @step(num_workers=4)\n",
    "    async def step_two(self, ctx: Context, ev: StepTwoEvent) -> StepThreeEvent:\n",
    "        print(\"Running query \", ev.query)\n",
    "        await asyncio.sleep(random.randint(1, 5))\n",
    "        return StepThreeEvent(result=ev.query)\n",
    "\n",
    "    @step\n",
    "    async def step_three(self, ctx: Context, ev: StepThreeEvent) -> StopEvent:\n",
    "        # wait until we receive 3 events\n",
    "        result = ctx.collect_events(ev, [StepThreeEvent] * 3)\n",
    "        if result is None:\n",
    "            print(\"Not all events received yet.\")\n",
    "            return None\n",
    "\n",
    "        # do something with all 3 results together\n",
    "        print(result)\n",
    "        return StopEvent(result=\"Done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1048887e-1503-4c64-ad72-ceb0def6ec25",
   "metadata": {
    "height": 64
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running query  Query 1\n",
      "Running query  Query 2\n",
      "Running query  Query 3\n",
      "Not all events received yet.\n",
      "Not all events received yet.\n",
      "[StepThreeEvent(result='Query 1'), StepThreeEvent(result='Query 3'), StepThreeEvent(result='Query 2')]\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "w = ConcurrentFlow(timeout=10, verbose=False)\n",
    "result = await w.run(message=\"Start the workflow.\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7586cbc4-c7b3-4f67-a245-38a274ae5cde",
   "metadata": {},
   "source": [
    "What `collect_events` does is store the events in the context until it has collected the number and type of events specified in its second argument. In this case, you've told it to wait for 3 events.\n",
    "\n",
    "If an event fires and `collect_events` hasn't yet seen the right number of events, it returns `None`, so you tell `step_three` to do nothing in that case. When `collect_events` receives the right number of events it returns them as an array, which you can see being printed in the final output.\n",
    "\n",
    "*Note:* This flow control lets you perform map-reduce style tasks.To implement a map-reduce pattern, you would split your task up into as many steps as necessary, and use `Context` to store that number with `ctx.set(\"num_events\", some_number)`. Then in `step_three` you would wait for the number stored in the context using `await ctx.get(\"num_events\")`. So you don't need to know in advance exactly how many concurrent steps you're taking. You'll do exactly this in a later lesson."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb6af0be-bedc-4efd-bd33-666f147b7590",
   "metadata": {},
   "source": [
    "### Collecting different event types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6033c923-8fd6-473e-b771-356fa6676b29",
   "metadata": {},
   "source": [
    "You don't just have to wait for multiple events of the same kind. In this example, you'll emit 3 totally different events and collect them at the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "84dc46d2-7b98-4720-b0c3-5b7b456a48cf",
   "metadata": {
    "height": 302
   },
   "outputs": [],
   "source": [
    "class StepAEvent(Event):\n",
    "    query: str\n",
    "\n",
    "class StepACompleteEvent(Event):\n",
    "    result: str\n",
    "\n",
    "class StepBEvent(Event):\n",
    "    query: str\n",
    "\n",
    "class StepBCompleteEvent(Event):\n",
    "    result: str\n",
    "\n",
    "class StepCEvent(Event):\n",
    "    query: str\n",
    "\n",
    "class StepCCompleteEvent(Event):\n",
    "    result: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c524727b-1d21-4aae-a551-ec115d94d684",
   "metadata": {
    "height": 759
   },
   "outputs": [],
   "source": [
    "class ConcurrentFlow(Workflow):\n",
    "    @step\n",
    "    async def start(\n",
    "        self, ctx: Context, ev: StartEvent\n",
    "    ) -> StepAEvent | StepBEvent | StepCEvent:\n",
    "        ctx.send_event(StepAEvent(query=\"Query 1\"))\n",
    "        ctx.send_event(StepBEvent(query=\"Query 2\"))\n",
    "        ctx.send_event(StepCEvent(query=\"Query 3\"))\n",
    "\n",
    "    @step\n",
    "    async def step_a(self, ctx: Context, ev: StepAEvent) -> StepACompleteEvent:\n",
    "        print(\"Doing something A-ish\")\n",
    "        return StepACompleteEvent(result=ev.query)\n",
    "\n",
    "    @step\n",
    "    async def step_b(self, ctx: Context, ev: StepBEvent) -> StepBCompleteEvent:\n",
    "        print(\"Doing something B-ish\")\n",
    "        return StepBCompleteEvent(result=ev.query)\n",
    "\n",
    "    @step\n",
    "    async def step_c(self, ctx: Context, ev: StepCEvent) -> StepCCompleteEvent:\n",
    "        print(\"Doing something C-ish\")\n",
    "        return StepCCompleteEvent(result=ev.query)\n",
    "\n",
    "    @step\n",
    "    async def step_three(\n",
    "        self,\n",
    "        ctx: Context,\n",
    "        ev: StepACompleteEvent | StepBCompleteEvent | StepCCompleteEvent,\n",
    "    ) -> StopEvent:\n",
    "        print(\"Received event \", ev.result)\n",
    "\n",
    "        # wait until we receive 3 events\n",
    "        events = ctx.collect_events(\n",
    "            ev,\n",
    "            [StepCCompleteEvent, StepACompleteEvent, StepBCompleteEvent],\n",
    "        )\n",
    "        if (events is None):\n",
    "            return None\n",
    "\n",
    "        # do something with all 3 results together\n",
    "        print(\"All events received: \", events)\n",
    "        return StopEvent(result=\"Done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61901db5-817c-4a75-b0cf-4d1eeab62bc6",
   "metadata": {},
   "source": [
    "When you run it, it will do all three things and wait for them in `step_three`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2e36f5c5-1057-47f0-b072-7b82c2478c8b",
   "metadata": {
    "height": 64
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Doing something A-ish\n",
      "Doing something B-ish\n",
      "Doing something C-ish\n",
      "Received event  Query 1\n",
      "Received event  Query 2\n",
      "Received event  Query 3\n",
      "All events received:  [StepCCompleteEvent(result='Query 3'), StepACompleteEvent(result='Query 1'), StepBCompleteEvent(result='Query 2')]\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "w = ConcurrentFlow(timeout=10, verbose=False)\n",
    "result = await w.run(message=\"Start the workflow.\")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9702b7bd-7f69-4ef6-a8d2-57460c201693",
   "metadata": {},
   "source": [
    "This new flow has quite a pretty visualization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8c893249-c71b-4597-adce-04dea538bf9b",
   "metadata": {
    "height": 79
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'NoneType'>\n",
      "<class '__main__.StepAEvent'>\n",
      "<class '__main__.StepBEvent'>\n",
      "<class '__main__.StepCEvent'>\n",
      "<class '__main__.StepACompleteEvent'>\n",
      "<class '__main__.StepBCompleteEvent'>\n",
      "<class '__main__.StepCCompleteEvent'>\n",
      "<class 'llama_index.core.workflow.events.StopEvent'>\n",
      "workflows/concurrent_different_events.html\n"
     ]
    }
   ],
   "source": [
    "WORKFLOW_FILE = \"workflows/concurrent_different_events.html\"\n",
    "draw_all_possible_flows(w, filename=WORKFLOW_FILE)\n",
    "html_content = extract_html_content(WORKFLOW_FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "7822a1d9-febd-4125-8f14-c4a5965a2aa7",
   "metadata": {
    "height": 45
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       " <div style=\"width: 100%; height: 800px; overflow: hidden;\"> <html>\n",
       "    <head>\n",
       "        <meta charset=\"utf-8\">\n",
       "        \n",
       "            <script src=\"lib/bindings/utils.js\"></script>\n",
       "            <link rel=\"stylesheet\" href=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/dist/vis-network.min.css\" integrity=\"sha512-WgxfT5LWjfszlPHXRmBWHkV2eceiWTOBvrKCNbdgDYTHrT2AeLCGbF4sZlZw3UMN3WtL0tGUoIAKsu8mllg/XA==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\" />\n",
       "            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js\" integrity=\"sha512-LnvoEWDFrqGHlHmDD2101OrLcbsfkrzoSpvtSQtxK3RMnRV0eOkhhBN2dXHKRrUU8p2DGRTk35n4O8nWSVe1mQ==\" crossorigin=\"anonymous\" referrerpolicy=\"no-referrer\"></script>\n",
       "            \n",
       "        \n",
       "<center>\n",
       "<h1></h1>\n",
       "</center>\n",
       "\n",
       "<!-- <link rel=\"stylesheet\" href=\"../node_modules/vis/dist/vis.min.css\" type=\"text/css\" />\n",
       "<script type=\"text/javascript\" src=\"../node_modules/vis/dist/vis.js\"> </script>-->\n",
       "        <link\n",
       "          href=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/css/bootstrap.min.css\"\n",
       "          rel=\"stylesheet\"\n",
       "          integrity=\"sha384-eOJMYsd53ii+scO/bJGFsiCZc+5NDVN2yr8+0RDqr0Ql0h+rP48ckxlpbzKgwra6\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        />\n",
       "        <script\n",
       "          src=\"https://cdn.jsdelivr.net/npm/bootstrap@5.0.0-beta3/dist/js/bootstrap.bundle.min.js\"\n",
       "          integrity=\"sha384-JEW9xMcG8R+pH31jmWH6WWP0WintQrMb4s7ZOdauHnUtxwoG2vI5DkLtS3qm9Ekf\"\n",
       "          crossorigin=\"anonymous\"\n",
       "        ></script>\n",
       "\n",
       "\n",
       "        <center>\n",
       "          <h1></h1>\n",
       "        </center>\n",
       "        <style type=\"text/css\">\n",
       "\n",
       "             #mynetwork {\n",
       "                 width: 100%;\n",
       "                 height: 750px;\n",
       "                 background-color: #ffffff;\n",
       "                 border: 1px solid lightgray;\n",
       "                 position: relative;\n",
       "                 float: left;\n",
       "             }\n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "\n",
       "             \n",
       "        </style>\n",
       "    </head>\n",
       "\n",
       "\n",
       "    <body>\n",
       "        <div class=\"card\" style=\"width: 100%\">\n",
       "            \n",
       "            \n",
       "            <div id=\"mynetwork\" class=\"card-body\"></div>\n",
       "        </div>\n",
       "\n",
       "        \n",
       "        \n",
       "\n",
       "        <script type=\"text/javascript\">\n",
       "\n",
       "              // initialize global variables.\n",
       "              var edges;\n",
       "              var nodes;\n",
       "              var allNodes;\n",
       "              var allEdges;\n",
       "              var nodeColors;\n",
       "              var originalNodes;\n",
       "              var network;\n",
       "              var container;\n",
       "              var options, data;\n",
       "              var filter = {\n",
       "                  item : '',\n",
       "                  property : '',\n",
       "                  value : []\n",
       "              };\n",
       "\n",
       "              \n",
       "\n",
       "              \n",
       "\n",
       "              // This method is responsible for drawing the graph, returns the drawn network\n",
       "              function drawGraph() {\n",
       "                  var container = document.getElementById('mynetwork');\n",
       "\n",
       "                  \n",
       "\n",
       "                  // parsing and collecting nodes and edges from the python\n",
       "                  nodes = new vis.DataSet([{\"color\": \"#FFA07A\", \"id\": \"StopEvent\", \"label\": \"StopEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"_done\", \"label\": \"_done\", \"shape\": \"box\"}, {\"color\": \"#ADD8E6\", \"id\": \"start\", \"label\": \"start\", \"shape\": \"box\"}, {\"color\": \"#E27AFF\", \"id\": \"StartEvent\", \"label\": \"StartEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_a\", \"label\": \"step_a\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"StepAEvent\", \"label\": \"StepAEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_b\", \"label\": \"step_b\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"StepBEvent\", \"label\": \"StepBEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_c\", \"label\": \"step_c\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"StepCEvent\", \"label\": \"StepCEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#ADD8E6\", \"id\": \"step_three\", \"label\": \"step_three\", \"shape\": \"box\"}, {\"color\": \"#90EE90\", \"id\": \"StepACompleteEvent\", \"label\": \"StepACompleteEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#90EE90\", \"id\": \"StepBCompleteEvent\", \"label\": \"StepBCompleteEvent\", \"shape\": \"ellipse\"}, {\"color\": \"#90EE90\", \"id\": \"StepCCompleteEvent\", \"label\": \"StepCCompleteEvent\", \"shape\": \"ellipse\"}]);\n",
       "                  edges = new vis.DataSet([{\"arrows\": \"to\", \"from\": \"StopEvent\", \"to\": \"_done\"}, {\"arrows\": \"to\", \"from\": \"start\", \"to\": \"StepAEvent\"}, {\"arrows\": \"to\", \"from\": \"start\", \"to\": \"StepBEvent\"}, {\"arrows\": \"to\", \"from\": \"start\", \"to\": \"StepCEvent\"}, {\"arrows\": \"to\", \"from\": \"StartEvent\", \"to\": \"start\"}, {\"arrows\": \"to\", \"from\": \"step_a\", \"to\": \"StepACompleteEvent\"}, {\"arrows\": \"to\", \"from\": \"StepAEvent\", \"to\": \"step_a\"}, {\"arrows\": \"to\", \"from\": \"step_b\", \"to\": \"StepBCompleteEvent\"}, {\"arrows\": \"to\", \"from\": \"StepBEvent\", \"to\": \"step_b\"}, {\"arrows\": \"to\", \"from\": \"step_c\", \"to\": \"StepCCompleteEvent\"}, {\"arrows\": \"to\", \"from\": \"StepCEvent\", \"to\": \"step_c\"}, {\"arrows\": \"to\", \"from\": \"step_three\", \"to\": \"StopEvent\"}, {\"arrows\": \"to\", \"from\": \"StepACompleteEvent\", \"to\": \"step_three\"}, {\"arrows\": \"to\", \"from\": \"StepBCompleteEvent\", \"to\": \"step_three\"}, {\"arrows\": \"to\", \"from\": \"StepCCompleteEvent\", \"to\": \"step_three\"}]);\n",
       "\n",
       "                  nodeColors = {};\n",
       "                  allNodes = nodes.get({ returnType: \"Object\" });\n",
       "                  for (nodeId in allNodes) {\n",
       "                    nodeColors[nodeId] = allNodes[nodeId].color;\n",
       "                  }\n",
       "                  allEdges = edges.get({ returnType: \"Object\" });\n",
       "                  // adding nodes and edges to the graph\n",
       "                  data = {nodes: nodes, edges: edges};\n",
       "\n",
       "                  var options = {\n",
       "    \"configure\": {\n",
       "        \"enabled\": false\n",
       "    },\n",
       "    \"edges\": {\n",
       "        \"color\": {\n",
       "            \"inherit\": true\n",
       "        },\n",
       "        \"smooth\": {\n",
       "            \"enabled\": true,\n",
       "            \"type\": \"dynamic\"\n",
       "        }\n",
       "    },\n",
       "    \"interaction\": {\n",
       "        \"dragNodes\": true,\n",
       "        \"hideEdgesOnDrag\": false,\n",
       "        \"hideNodesOnDrag\": false\n",
       "    },\n",
       "    \"physics\": {\n",
       "        \"enabled\": true,\n",
       "        \"stabilization\": {\n",
       "            \"enabled\": true,\n",
       "            \"fit\": true,\n",
       "            \"iterations\": 1000,\n",
       "            \"onlyDynamicEdges\": false,\n",
       "            \"updateInterval\": 50\n",
       "        }\n",
       "    }\n",
       "};\n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  network = new vis.Network(container, data, options);\n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "                  \n",
       "\n",
       "\n",
       "                  \n",
       "\n",
       "                  return network;\n",
       "\n",
       "              }\n",
       "              drawGraph();\n",
       "        </script>\n",
       "    </body>\n",
       "</html> </div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {
      "isolated": true
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(HTML(html_content), metadata=dict(isolated=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5927ba40-53fc-4710-98f2-359928e98f46",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d12e619-fbed-4fcd-8a12-f16147f16a56",
   "metadata": {},
   "source": [
    "In practical use, agents can take a long time to run. It's a poor user-experience to have the user execute a workflow and then wait a long time to see if it works or not; it's better to give them some indication that things are happening in real-time, even if the process is not complete.\n",
    "\n",
    "To do this, Workflows allow streaming events back to the user. Here you'll use `Context.write_event_to_stream` to emit these events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "775c33ea-fcaa-48b1-afe4-0220db6b9bfe",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b32e825f-4539-4992-8b78-bde29ddd1c7a",
   "metadata": {
    "height": 217
   },
   "outputs": [],
   "source": [
    "class FirstEvent(Event):\n",
    "    first_output: str\n",
    "\n",
    "class SecondEvent(Event):\n",
    "    second_output: str\n",
    "    response: str\n",
    "\n",
    "class TextEvent(Event):\n",
    "    delta: str\n",
    "\n",
    "class ProgressEvent(Event):\n",
    "    msg: str"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a2079c9-e7eb-4e84-a135-a301a4a016c5",
   "metadata": {},
   "source": [
    "The specific event we'll be sending back is the \"delta\" responses from the LLM. When you ask an LLM to generate a streaming response as you're doing here, it sends back each chunk of its response as it becomes available. This is available as the \"delta\". You're going to wrap the delta in a TextEvent and send it back to the Workflow's own stream."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "d0e9d392-8a2d-4411-aef2-a5e93102d5d5",
   "metadata": {
    "height": 436
   },
   "outputs": [],
   "source": [
    "class MyWorkflow(Workflow):\n",
    "    @step\n",
    "    async def step_one(self, ctx: Context, ev: StartEvent) -> FirstEvent:\n",
    "        ctx.write_event_to_stream(ProgressEvent(msg=\"Step one is happening\"))\n",
    "        return FirstEvent(first_output=\"First step complete.\")\n",
    "\n",
    "    @step\n",
    "    async def step_two(self, ctx: Context, ev: FirstEvent) -> SecondEvent:\n",
    "        llm = OpenAI(model=\"gpt-4o-mini\", api_key=api_key) \n",
    "        generator = await llm.astream_complete(\n",
    "            \"Please give me the first 50 words of Moby Dick, a book in the public domain.\"\n",
    "        )\n",
    "        async for response in generator:\n",
    "            # Allow the workflow to stream this piece of response\n",
    "            ctx.write_event_to_stream(TextEvent(delta=response.delta))\n",
    "        return SecondEvent(\n",
    "            second_output=\"Second step complete, full response attached\",\n",
    "            response=str(response),\n",
    "        )\n",
    "\n",
    "    @step\n",
    "    async def step_three(self, ctx: Context, ev: SecondEvent) -> StopEvent:\n",
    "        ctx.write_event_to_stream(ProgressEvent(msg=\"Step three is happening\"))\n",
    "        return StopEvent(result=\"Workflow complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be37ded4-8587-406a-b730-4082f573a6c2",
   "metadata": {},
   "source": [
    "You can work with the emitted events by getting a streaming endpoint from the `run` command, and then filtering it for the types of events we want to see (you could print every event if you wanted to but that would be quite noisy).\n",
    "\n",
    "In this case you'll just print out the progressevents and the textevents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "10eaae2e-6950-4192-930d-a1c7cca483c8",
   "metadata": {
    "height": 215
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step one is happening\n",
      "Sure! Here are the first 50 words of \"Moby Dick\" by Herman Melville:\n",
      "\n",
      "\"Call me Ishmael. Some years ago—never mind how long precisely—I had little or no money in my purse, and nothing particular to interest me on shore. So, whenever I find myself growing grim about the mouth; whenever it is a damp, drizzly November in my soul;\"Step three is happening\n",
      "Final result =  Workflow complete.\n"
     ]
    }
   ],
   "source": [
    "workflow = MyWorkflow(timeout=30, verbose=False)\n",
    "handler = workflow.run(first_input=\"Start the workflow.\")\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, ProgressEvent):\n",
    "        print(ev.msg)\n",
    "    if isinstance(ev, TextEvent):\n",
    "        print(ev.delta, end=\"\")\n",
    "\n",
    "final_result = await handler\n",
    "print(\"Final result = \", final_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8d81176-340f-4dd6-89be-97296d46023e",
   "metadata": {},
   "source": [
    "## Congratulations!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87272551-d5a1-4998-af3d-696fa660bbff",
   "metadata": {},
   "source": [
    "You've successfully built a number of increasingly-complex Workflows, and mastered the basic concepts. In the next lesson, you'll add RAG to your Workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36a004fa",
   "metadata": {},
   "source": [
    "## Resource"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d1f43b8",
   "metadata": {},
   "source": [
    "[Guide to LlamaIndex's Workflows](https://docs.llamaindex.ai/en/stable/module_guides/workflow/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}