Fixed "LoopEvent" handling in Loops and Branches

unit2/llama-index/workflows.ipynb

@@ -146,13 +146,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "Bad thing happened\n",
       "Good thing happened\n"
      ]
     },
@@ -182,7 +183,7 @@
     "\n",
     "class MultiStepWorkflow(Workflow):\n",
     "    @step\n",
-    "    async def step_one(self, ev: StartEvent) -> ProcessingEvent | LoopEvent:\n",
+    "    async def step_one(self, ev: StartEvent | LoopEvent) -> ProcessingEvent | LoopEvent:\n",
     "        if random.randint(0, 1) == 0:\n",
     "            print(\"Bad thing happened\")\n",
     "            return LoopEvent(loop_output=\"Back to step one.\")\n",
@@ -191,7 +192,7 @@
     "            return ProcessingEvent(intermediate_result=\"First step complete.\")\n",
     "\n",
     "    @step\n",
-    "    async def step_two(self, ev: ProcessingEvent | LoopEvent) -> StopEvent:\n",
+    "    async def step_two(self, ev: ProcessingEvent) -> StopEvent:\n",
     "        # Use the intermediate result\n",
     "        final_result = f\"Finished processing: {ev.intermediate_result}\"\n",
     "        return StopEvent(result=final_result)\n",