From d7dc16310fc598c92e16d5b31adebce7a6105920 Mon Sep 17 00:00:00 2001
From: isaac hershenson <ihershenson@hmc.edu>
Date: Fri, 19 Jul 2024 13:05:28 -0700
Subject: [PATCH] nuno comments

---
 examples/node-retries.ipynb | 76 ++++++++++++-------------------------
 1 file changed, 25 insertions(+), 51 deletions(-)

diff --git a/examples/node-retries.ipynb b/examples/node-retries.ipynb
index 29f59e7dd..93749e065 100644
--- a/examples/node-retries.ipynb
+++ b/examples/node-retries.ipynb
@@ -6,113 +6,87 @@
    "source": [
     "# How to add node retry policies\n",
     "\n",
-    "There are many use cases where you may wish for your node to have a custom retry policy, for example if you are making HTTP requests, calling an API, etc. \n",
+    "There are many use cases where you may wish for your node to have a custom retry policy, for example if you are calling an API, querying a databse, or calling an LLM, etc. \n",
     "\n",
-    "## Defining your retry policy\n",
-    "\n",
-    "In order to configure the retry policty, you have to pass the `retry` parameter to the `add_node` function. The `retry` parameter takes in a `RetryPolicy` named tuple object. Below we instantiate a `RetryPolicy` object with the default parameters."
+    "In order to configure the retry policty, you have to pass the `retry` parameter to the `add_node` function. The `retry` parameter takes in a `RetryPolicy` named tuple object. Below we instantiate a `RetryPolicy` object with the default parameters:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from langgraph.pregel.types import default_retry_on\n",
-    "from langgraph.pregel import RetryPolicy\n",
-    "\n",
-    "# Note: This is equivalent to writing retry_policy=RetryPolicy() since these are all default values\n",
-    "retry_policy = RetryPolicy(\n",
-    "    initial_interval = 0.5,\n",
-    "    backoff_factor = 2.0,\n",
-    "    max_interval = 128.0,\n",
-    "    max_attempts = 3,\n",
-    "    jitter = True,\n",
-    "    retry_on = default_retry_on\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "You can then change the default values to work for your custom use case:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "RetryPolicy(initial_interval=0.5, backoff_factor=2.0, max_interval=128.0, max_attempts=5, jitter=True, retry_on=<function default_retry_on at 0x1157419e0>)"
+       "RetryPolicy(initial_interval=0.5, backoff_factor=2.0, max_interval=128.0, max_attempts=3, jitter=True, retry_on=<function default_retry_on at 0x1157419e0>)"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "retry_policy = retry_policy._replace(max_attempts = 5)\n",
-    "retry_policy"
+    "from langgraph.pregel import RetryPolicy\n",
+    "\n",
+    "RetryPolicy()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "If you want more information on what each of the parameters does, be sure to read the source code.\n",
+    "If you want more information on what each of the parameters does, be sure to read the [reference](https://langchain-ai.github.io/langgraph/reference/graphs/#langgraph.graph.StateGraph.add_node).\n",
     "\n",
-    "## Passing your retry policy to a node\n",
+    "## Passing a retry policy to a node\n",
     "\n",
-    "Lastly, we can pass `retry_policy` to the `add_node` function:"
+    "Lastly, we can pass `RetryPolicy` objects when we call the `add_node` function. In the example below we pass two different retry policies to each of our nodes:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
     "import operator\n",
+    "import requests\n",
+    "import sqlite3\n",
     "from typing import Annotated, Sequence, TypedDict\n",
     "\n",
     "from langchain_anthropic import ChatAnthropic\n",
     "from langchain_core.messages import BaseMessage\n",
     "\n",
     "from langgraph.graph import END, StateGraph, START\n",
+    "from langchain_community.utilities import SQLDatabase\n",
+    "from langchain_core.messages import AIMessage\n",
     "\n",
+    "db = SQLDatabase.from_uri(\"sqlite:///:memory:\")\n",
     "\n",
     "model = ChatAnthropic(model_name=\"claude-2.1\")\n",
     "\n",
-    "\n",
     "class AgentState(TypedDict):\n",
     "    messages: Annotated[Sequence[BaseMessage], operator.add]\n",
     "\n",
+    "def query_database(state):\n",
+    "    query_result = db.run(\"SELECT * FROM Artist LIMIT 10;\")\n",
+    "    return {\"messages\": [AIMessage(content=query_result)]}\n",
     "\n",
-    "def _call_model(state):\n",
+    "def call_model(state):\n",
     "    response = model.invoke(state[\"messages\"])\n",
     "    return {\"messages\": [response]}\n",
     "\n",
     "# Define a new graph\n",
     "workflow = StateGraph(AgentState)\n",
-    "workflow.add_node(\"model\", _call_model,retry=retry_policy)\n",
+    "workflow.add_node(\"query_database\",query_database, retry=RetryPolicy(retry_on=sqlite3.OperationalError))\n",
+    "workflow.add_node(\"model\", call_model, retry=RetryPolicy(max_attempts=5))\n",
     "workflow.add_edge(START, \"model\")\n",
-    "workflow.add_edge(\"model\", END)\n",
+    "workflow.add_edge(\"model\",\"query_database\")\n",
+    "workflow.add_edge(\"query_database\", END)\n",
     "\n",
     "app = workflow.compile()"
    ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now the `\"model\"` node will be configured to use the custom retry policy that was defined above."
-   ]
   }
  ],
  "metadata": {