diff --git a/docs/_scripts/prepare_notebooks_for_ci.py b/docs/_scripts/prepare_notebooks_for_ci.py
index fe038fa67..83267c3c6 100644
--- a/docs/_scripts/prepare_notebooks_for_ci.py
+++ b/docs/_scripts/prepare_notebooks_for_ci.py
@@ -88,10 +88,11 @@ def has_blocklisted_command(code: str, metadata: dict) -> bool:
             return True
     return False
 
-def add_mermaid_retries(code: str) -> str:
+def remove_mermaid(code: str) -> str:
     return code.replace(
-        "draw_mermaid_png()",
-        "draw_mermaid_png(max_retries=10, retry_delay=2.0)"
+        "display(Image(graph.get_graph().draw_mermaid_png()))",
+        # replace with a dummy statement
+        "print()"
     )
 
 
@@ -188,12 +189,12 @@ def add_vcr_to_notebook(
     return notebook
 
 
-def add_mermaid_retries_to_notebook(notebook: nbformat.NotebookNode) -> nbformat.NotebookNode:
+def remove_mermaid_from_notebook(notebook: nbformat.NotebookNode) -> nbformat.NotebookNode:
     for cell in notebook.cells:
         if cell.cell_type != "code":
             continue
 
-        cell.source = add_mermaid_retries(cell.source)
+        cell.source = remove_mermaid(cell.source)
     return notebook
 
 
@@ -218,7 +219,7 @@ def process_notebooks(should_comment_install_cells: bool) -> None:
                             notebook, cassette_prefix=cassette_prefix
                         )
 
-                    notebook = add_mermaid_retries_to_notebook(notebook)
+                    notebook = remove_mermaid_from_notebook(notebook)
 
                     if notebook_path in NOTEBOOKS_NO_EXECUTION:
                         # Add a cell at the beginning to indicate that this notebook should not be executed
diff --git a/docs/docs/how-tos/human_in_the_loop/wait-user-input.ipynb b/docs/docs/how-tos/human_in_the_loop/wait-user-input.ipynb
index 6ffca2b6b..770d8bf4a 100644
--- a/docs/docs/how-tos/human_in_the_loop/wait-user-input.ipynb
+++ b/docs/docs/how-tos/human_in_the_loop/wait-user-input.ipynb
@@ -499,7 +499,8 @@
     "    config,\n",
     "    stream_mode=\"values\",\n",
     "):\n",
-    "    event[\"messages\"][-1].pretty_print()"
+    "    if \"messages\" in event:\n",
+    "        event[\"messages\"][-1].pretty_print()"
    ]
   },
   {
@@ -572,7 +573,8 @@
    ],
    "source": [
     "for event in app.stream(Command(resume=\"san francisco\"), config, stream_mode=\"values\"):\n",
-    "    event[\"messages\"][-1].pretty_print()"
+    "    if \"messages\" in event:\n",
+    "        event[\"messages\"][-1].pretty_print()"
    ]
   }
  ],
@@ -592,7 +594,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,