{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 calling an API, querying a database, or calling an LLM, etc. \n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install the required packages and set our API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain_anthropic langchain_community"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"ANTHROPIC_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to configure the retry policy, 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": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "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": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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 [reference](https://langchain-ai.github.io/langgraph/reference/graphs/#retrypolicy).\n",
    "\n",
    "## Passing a retry policy to a node\n",
    "\n",
    "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": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\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",
    "\n",
    "def query_database(state):\n",
    "    query_result = db.run(\"SELECT * FROM Artist LIMIT 10;\")\n",
    "    return {\"messages\": [AIMessage(content=query_result)]}\n",
    "\n",
    "\n",
    "def call_model(state):\n",
    "    response = model.invoke(state[\"messages\"])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "workflow.add_node(\n",
    "    \"query_database\",\n",
    "    query_database,\n",
    "    retry=RetryPolicy(retry_on=sqlite3.OperationalError),\n",
    ")\n",
    "workflow.add_node(\"model\", call_model, retry=RetryPolicy(max_attempts=5))\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_edge(\"model\", \"query_database\")\n",
    "workflow.add_edge(\"query_database\", END)\n",
    "\n",
    "app = workflow.compile()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}