From 2c8f1a92e01c1aa15203fbdc877992127d525135 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Thu, 18 Apr 2024 23:16:12 -0700
Subject: [PATCH] Fix conversational RAG example for collab

An import wont work because google collab doesn't load
surrounding files.
---
 examples/conversational-rag/notebook.ipynb | 38 +++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/examples/conversational-rag/notebook.ipynb b/examples/conversational-rag/notebook.ipynb
index ca7dca3d6..bce70c668 100644
--- a/examples/conversational-rag/notebook.ipynb
+++ b/examples/conversational-rag/notebook.ipynb
@@ -47,7 +47,6 @@
     "import burr.core\n",
     "from burr.core import ApplicationBuilder, State, default, expr\n",
     "from burr.core.action import action\n",
-    "from application import PrintStepHook # local import from application.py\n",
     "from burr.tracking import LocalTrackingClient"
    ]
   },
@@ -335,6 +334,43 @@
     "    return {\"question\": user_question}, state"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Add a hook to print the steps -- optional but shows that Burr is pluggable"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "4439e0dd39441414"
+  },
+  {
+   "cell_type": "code",
+   "outputs": [],
+   "source": [
+    "from burr.core import Action\n",
+    "from burr.lifecycle import PostRunStepHook, PreRunStepHook\n",
+    "\n",
+    "class PrintStepHook(PostRunStepHook, PreRunStepHook):\n",
+    "    \"\"\"Custom hook to print the action/result after each step.\"\"\"\n",
+    "\n",
+    "    def pre_run_step(self, action: Action, **future_kwargs):\n",
+    "        if action.name == \"ai_converse\":\n",
+    "            print(\"🤔 AI is thinking...\")\n",
+    "        if action.name == \"human_converse\":\n",
+    "            print(\"⏳Processing input from user...\")\n",
+    "\n",
+    "    def post_run_step(self, *, state: \"State\", action: Action, result: dict, **future_kwargs):\n",
+    "        if action.name == \"human_converse\":\n",
+    "            print(\"🎙💬\", result[\"question\"], \"\\n\")\n",
+    "        if action.name == \"ai_converse\":\n",
+    "            print(\"🤖💬\", result[\"conversational_rag_response\"], \"\\n\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "e3c3af84fab14f39"
+  },
   {
    "cell_type": "markdown",
    "id": "9579b47aac2c53a0",