{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8bcd1a3d-7c50-4f58-be4e-1ed654aa33be",
   "metadata": {},
   "source": [
    "# How to visualize your graph\n",
    "\n",
    "This notebook walks through how to visualize the graphs you create. This works with ANY [Graph](https://langchain-ai.github.io/langgraph/reference/graphs/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32a0e7f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e130cf70-a30e-47d7-8fd5-464f1a92e374",
   "metadata": {},
   "source": [
    "## Set up Graph\n",
    "\n",
    "You can visualize any arbitrary Graph, including StateGraph's and MessageGraph's. Let's have some fun by drawing fractals :)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6d604311",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from typing import Annotated, Literal\n",
    "\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "class MyNode:\n",
    "    def __init__(self, name: str):\n",
    "        self.name = name\n",
    "\n",
    "    def __call__(self, state: State):\n",
    "        return {\"messages\": [(\"assistant\", f\"Called node {self.name}\")]}\n",
    "\n",
    "\n",
    "def route(state) -> Literal[\"entry_node\", \"__end__\"]:\n",
    "    if len(state[\"messages\"]) > 10:\n",
    "        return \"__end__\"\n",
    "    return \"entry_node\"\n",
    "\n",
    "\n",
    "def add_fractal_nodes(builder, current_node, level, max_level):\n",
    "    if level > max_level:\n",
    "        return\n",
    "\n",
    "    # Number of nodes to create at this level\n",
    "    num_nodes = random.randint(1, 3)  # Adjust randomness as needed\n",
    "    for i in range(num_nodes):\n",
    "        nm = [\"A\", \"B\", \"C\"][i]\n",
    "        node_name = f\"node_{current_node}_{nm}\"\n",
    "        builder.add_node(node_name, MyNode(node_name))\n",
    "        builder.add_edge(current_node, node_name)\n",
    "\n",
    "        # Recursively add more nodes\n",
    "        r = random.random()\n",
    "        if r > 0.2 and level + 1 < max_level:\n",
    "            add_fractal_nodes(builder, node_name, level + 1, max_level)\n",
    "        elif r > 0.05:\n",
    "            builder.add_conditional_edges(node_name, route, node_name)\n",
    "        else:\n",
    "            # End\n",
    "            builder.add_edge(node_name, \"__end__\")\n",
    "\n",
    "\n",
    "def build_fractal_graph(max_level: int):\n",
    "    builder = StateGraph(State)\n",
    "    entry_point = \"entry_node\"\n",
    "    builder.add_node(entry_point, MyNode(entry_point))\n",
    "    builder.add_edge(START, entry_point)\n",
    "\n",
    "    add_fractal_nodes(builder, entry_point, 1, max_level)\n",
    "\n",
    "    # Optional: set a finish point if required\n",
    "    builder.add_edge(entry_point, END)  # or any specific node\n",
    "\n",
    "    return builder.compile()\n",
    "\n",
    "\n",
    "app = build_fractal_graph(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4fc9378-b141-4b65-b86c-3afba77f7161",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-18T12:18:30.605220Z",
     "start_time": "2024-04-18T12:18:30.587191Z"
    }
   },
   "source": [
    "## Ascii\n",
    "\n",
    "We can easily visualize this graph in ascii"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ca9b980d-1f0a-4286-9157-a870e3d55134",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-19T11:25:37.303260Z",
     "start_time": "2024-04-19T11:25:37.273032Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                           +-----------+                                                               \n",
      "                                                                           | __start__ |                                                               \n",
      "                                                                           +-----------+                                                               \n",
      "                                                                                  *                                                                    \n",
      "                                                                                  *                                                                    \n",
      "                                                                                  *                                                                    \n",
      "                                                                          +------------+                                                               \n",
      "                                                                    ******| entry_node |..*****                                                        \n",
      "                                                        ************ *****+------------+  ......***********                                            \n",
      "                                           *************        *****             .             .....      ************                                \n",
      "                               ************               ******                  .                  .....             ************                    \n",
      "                        *******                      *****                        .                       ......                   ************        \n",
      "    +-------------------+                         ***                            ..                             ...                            ******* \n",
      "    | node_entry_node_B |*********                  **                        ...                                 .                                  * \n",
      "    +-------------------+         ******************* **                   ...                                    .                                  * \n",
      "              *                                      *******************...                                       .                                  * \n",
      "              *                                           **         ...*******************                       .                                  * \n",
      "              *                                             **     ..                      **********             .                                  * \n",
      "+--------------------------+                         +-------------------+                          +--------------------------+                ****** \n",
      "| node_node_entry_node_B_A |***                      | node_entry_node_A |                          | node_node_entry_node_B_B |      **********       \n",
      "+--------------------------+   **********            +-------------------+                          +--------------------------+******                 \n",
      "                                         **********                     ...                    .....             **********                            \n",
      "                                                   **********              ...           ......       ***********                                      \n",
      "                                                             **********       ..      ...   **********                                                 \n",
      "                                                                       *****+---------+*****                                                           \n",
      "                                                                            | __end__ |                                                                \n",
      "                                                                            +---------+                                                                \n"
     ]
    }
   ],
   "source": [
    "app.get_graph().print_ascii()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edcd9ad2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-18T12:18:30.629307Z",
     "start_time": "2024-04-18T12:18:30.609323Z"
    }
   },
   "source": [
    "## Mermaid\n",
    "\n",
    "We can also convert a graph class into Mermaid syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "66007b2d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-19T11:25:38.733126Z",
     "start_time": "2024-04-19T11:25:38.726838Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%%{init: {'flowchart': {'curve': 'linear'}}}%%\n",
      "graph TD;\n",
      "\t__start__[__start__]:::startclass;\n",
      "\t__end__[__end__]:::endclass;\n",
      "\tentry_node([entry_node]):::otherclass;\n",
      "\tnode_entry_node_A([node_entry_node_A]):::otherclass;\n",
      "\tnode_entry_node_B([node_entry_node_B]):::otherclass;\n",
      "\tnode_node_entry_node_B_A([node_node_entry_node_B_A]):::otherclass;\n",
      "\tnode_node_entry_node_B_B([node_node_entry_node_B_B]):::otherclass;\n",
      "\t__start__ --> entry_node;\n",
      "\tentry_node --> __end__;\n",
      "\tentry_node --> node_entry_node_A;\n",
      "\tentry_node --> node_entry_node_B;\n",
      "\tnode_entry_node_B --> node_node_entry_node_B_A;\n",
      "\tnode_entry_node_B --> node_node_entry_node_B_B;\n",
      "\tnode_node_entry_node_B_A --> __end__;\n",
      "\tnode_entry_node_A -.-> entry_node;\n",
      "\tnode_entry_node_A -.-> __end__;\n",
      "\tnode_node_entry_node_B_B -.-> entry_node;\n",
      "\tnode_node_entry_node_B_B -.-> __end__;\n",
      "\tclassDef startclass fill:#ffdfba;\n",
      "\tclassDef endclass fill:#baffc9;\n",
      "\tclassDef otherclass fill:#fad7de;\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(app.get_graph().draw_mermaid())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f77ad75",
   "metadata": {},
   "source": [
    "## PNG\n",
    "\n",
    "If preferred, we could render the Graph into a  `.png`. Here we could use three options:\n",
    "\n",
    "- Using Mermaid.ink API (does not require additional packages)\n",
    "- Using Mermaid + Pyppeteer (requires `pip install pyppeteer`)\n",
    "- Using graphviz (which requires `pip install graphviz`)\n",
    "\n",
    "\n",
    "### Using Mermaid.Ink\n",
    "\n",
    "By default, `draw_mermaid_png()` uses Mermaid.Ink's API to generate the diagram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "967f116d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 45364 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "from langchain_core.runnables.graph import CurveStyle, MermaidDrawMethod, NodeStyles\n",
    "\n",
    "display(\n",
    "    Image(\n",
    "        app.get_graph().draw_mermaid_png(\n",
    "            draw_method=MermaidDrawMethod.API,\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9e767fc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-18T12:18:30.873950Z",
     "start_time": "2024-04-18T12:18:30.871750Z"
    }
   },
   "source": [
    "### Using Mermaid + Pyppeteer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d403e1e7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-19T11:25:44.798703Z",
     "start_time": "2024-04-19T11:25:44.793438Z"
    }
   },
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --quiet pyppeteer\n",
    "%pip install --quiet nest_asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "058546ee",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-19T11:25:47.412695Z",
     "start_time": "2024-04-19T11:25:45.405158Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 166184 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()  # Required for Jupyter Notebook to run async functions\n",
    "\n",
    "display(\n",
    "    Image(\n",
    "        app.get_graph().draw_mermaid_png(\n",
    "            curve_style=CurveStyle.LINEAR,\n",
    "            node_colors=NodeStyles(first=\"#ffdfba\", last=\"#baffc9\", default=\"#fad7de\"),\n",
    "            wrap_label_n_words=9,\n",
    "            output_file_path=None,\n",
    "            draw_method=MermaidDrawMethod.PYPPETEER,\n",
    "            background_color=\"white\",\n",
    "            padding=10,\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d821b2f6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-18T12:18:30.629629Z",
     "start_time": "2024-04-18T12:18:30.620092Z"
    }
   },
   "source": [
    "### Using Graphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d4234400-75cd-4b13-aeff-828f7fb68ab1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-19T11:25:42.057704Z",
     "start_time": "2024-04-19T11:25:42.019017Z"
    }
   },
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install pygraphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ee026342-f560-4ce0-ab43-1718bd19a366",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-04-19T11:25:42.631675Z",
     "start_time": "2024-04-19T11:25:42.452377Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95372 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Image(app.get_graph().draw_png()))"
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
