Cell chaining in note books

IBM · Sep 26, 2024 · 8721f94 · 8721f94
1 parent da30bb0
commit 8721f94
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 42 deletions.
diff --git a/examples/notebooks/notebook.ipynb b/examples/notebooks/notebook.ipynb
@@ -113,7 +113,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6b72c204-b8de-45b2-9beb-9c6565d28de2",
+   "id": "91667b7f-e097-42cc-a1b2-765867ca6d40",
    "metadata": {},
    "outputs": [],
    "source": []

diff --git a/examples/notebooks/notebook_debug.ipynb b/examples/notebooks/notebook_debug.ipynb
@@ -16,6 +16,13 @@
    "id": "f3c62df1-0347-4711-acd7-3892cfd5df30",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Hello,\u001b[32m World\u001b[0m\u001b[32m!\u001b[0m"
+     ]
+    },
     {
      "data": {
       "text/html": [
@@ -50,7 +57,7 @@
        "      background-color: rgb(238, 184, 112);\n",
        "    }\n",
        "\n",
-       "    .pdl_document {\n",
+       "    .pdl_text {\n",
        "      background-color: rgb(219, 215, 250);\n",
        "    }\n",
        "\n",
@@ -136,7 +143,7 @@
        "          { type: 'main', style: pstyle, html: '<div id=\"code\"></div>' }\n",
        "        ]\n",
        "      })\n",
-       "      const data = {\"kind\": \"document\", \"description\": \"Model call\", \"defs\": {}, \"document\": [\"Hello,\", {\"kind\": \"model\", \"defs\": {}, \"platform\": \"litellm\", \"model\": \"watsonx/ibm/granite-34b-code-instruct\", \"parameters\": {\"stop\": [\"!\"], \"include_stop_sequence\": true}, \"result\": \" World!\", \"location\": {\"path\": [\"document\", \"[1]\"], \"file\": \"\", \"table\": {\"['description']\": 1, \"['document']\": 2, \"['document', '[0]']\": 3, \"['document', '[1]']\": 4, \"['document', '[1]', 'model']\": 4, \"['document', '[1]', 'parameters']\": 5, \"['document', '[1]', 'parameters', 'stop']\": 6, \"['document', '[1]', 'parameters', 'include_stop_sequence']\": 7}}}], \"result\": \"Hello, World!\", \"location\": {\"path\": [], \"file\": \"\", \"table\": {\"['description']\": 1, \"['document']\": 2, \"['document', '[0]']\": 3, \"['document', '[1]']\": 4, \"['document', '[1]', 'model']\": 4, \"['document', '[1]', 'parameters']\": 5, \"['document', '[1]', 'parameters', 'stop']\": 6, \"['document', '[1]', 'parameters', 'include_stop_sequence']\": 7}}}\n",
+       "      const data = {\"kind\": \"text\", \"description\": \"Model call\", \"defs\": {}, \"text\": [\"Hello,\", {\"kind\": \"model\", \"defs\": {}, \"platform\": \"litellm\", \"model\": \"watsonx/ibm/granite-34b-code-instruct\", \"parameters\": {\"stop\": [\"!\"], \"include_stop_sequence\": true}, \"result\": \" World!\", \"location\": {\"path\": [\"text\", \"[1]\"], \"file\": \"\", \"table\": {\"['description']\": 1, \"['text']\": 2, \"['text', '[0]']\": 3, \"['text', '[1]']\": 4, \"['text', '[1]', 'model']\": 4, \"['text', '[1]', 'parameters']\": 5, \"['text', '[1]', 'parameters', 'stop']\": 6, \"['text', '[1]', 'parameters', 'include_stop_sequence']\": 7}}}], \"result\": \"Hello, World!\", \"location\": {\"path\": [], \"file\": \"\", \"table\": {\"['description']\": 1, \"['text']\": 2, \"['text', '[0]']\": 3, \"['text', '[1]']\": 4, \"['text', '[1]', 'model']\": 4, \"['text', '[1]', 'parameters']\": 5, \"['text', '[1]', 'parameters', 'stop']\": 6, \"['text', '[1]', 'parameters', 'include_stop_sequence']\": 7}}}\n",
        "      pdl_viewer.replace_div('doc', pdl_viewer.show_output(data))\n",
        "\n",
        "    </script>\n",
@@ -153,7 +160,7 @@
     }
    ],
    "source": [
-    "%%pdl_debug\n",
+    "%%pdl --viewer\n",
     "description: Model call\n",
     "text: \n",
     "- Hello,\n",

diff --git a/pdl/pdl_notebook_ext.py b/pdl/pdl_notebook_ext.py
@@ -13,57 +13,34 @@
 class PDLMagics(Magics):
     @cell_magic
     @magic_arguments()
-    @argument(
-        "result_name",
-        nargs="?",
-        default=None,
-        help="""Variable where to store the result of the execution of the provided PDL program.""",
-    )
+    @argument('-r', '--reset-context', action='store_true',
+        default=False, help='Reset the background context to the empty list.')
+    @argument('--viewer', action='store_true',
+        default=False, help='Show the execution trace in the PDL viewer.')
     @needs_local_scope
     def pdl(self, line, cell, local_ns):
         line = line.strip()
         args = parse_argstring(self.pdl, line)
-        try:
-            result = exec_str(
-                cell,
-                config=InterpreterConfig(
-                    yield_result=True, yield_background=False, batch=0
-                ),
-                scope=local_ns,
-            )
-        except Exception as err:
-            print(err)
-            return
-        if args.result_name is not None:
-            local_ns[args.result_name] = result
-
-    @cell_magic
-    @magic_arguments()
-    @argument(
-        "result_name",
-        nargs="?",
-        default=None,
-        help="""Variable where to store the result of the execution of the provided PDL program.""",
-    )
-    @needs_local_scope
-    def pdl_debug(self, line, cell, local_ns):
-        line = line.strip()
-        args = parse_argstring(self.pdl_debug, line)
+        if args.reset_context:
+            scope = local_ns | {"context": []}
+        else:
+            scope = local_ns
         try:
             pdl_output = exec_str(
                 cell,
                 config=InterpreterConfig(
-                    yield_result=False, yield_background=False, batch=0
+                    yield_result=True, yield_background=False, batch=0
                 ),
-                scope=local_ns,
-                output="all",
+                scope=scope,
+                output="all"
             )
         except Exception as err:
             print(err)
             return
-        if args.result_name is not None:
-            local_ns[args.result_name] = pdl_output["result"]
-        display_html(self.pdl_viewer(block_to_dict(pdl_output["trace"])))
+        for x, v in pdl_output["scope"].items():
+            local_ns[x] = v
+        if args.viewer:
+          display_html(self.pdl_viewer(block_to_dict(pdl_output["trace"])))
 
     def pdl_viewer(self, trace):
         trace_str = json.dumps(trace)