{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ng45O5Y80vzt"
   },
   "source": [
    "# Google AlloyDB for PostgreSQL\n",
    "\n",
    "> [AlloyDB](https://cloud.google.com/alloydb) is a fully managed PostgreSQL compatible database service for your most demanding enterprise workloads. AlloyDB combines the best of Google with PostgreSQL, for superior performance, scale, and availability. Extend your database application to build AI-powered experiences leveraging AlloyDB Langgraph integrations.\n",
    "\n",
    "This notebook goes over how to use `AlloyDB for PostgreSQL` to store checkpoints with the `AlloyDBSaver` class.\n",
    "\n",
    "Learn more about the package on [GitHub](https://github.com/googleapis/langchain-google-alloydb-pg-python/).\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/googleapis/langchain-google-alloydb-pg-python/blob/langgraph-base/docs/langgraph_checkpoint.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Prerequisites\n",
    "\n",
    "This guide assumes familiarity with the following:\n",
    "\n",
    "- [LangGraph Persistence](https://langchain-ai.github.io/langgraph/concepts/persistence/)\n",
    "- [Postgresql](https://www.postgresql.org/about/)\n",
    "- [AlloyDB](https://cloud.google.com/products/alloydb?hl=en#how-it-works)\n",
    "\n",
    "When creating LangGraph agents, you can also set them up so that they persist their state. This allows you to do things like interact with an agent multiple times and have it remember previous interactions."
   ],
   "metadata": {
    "id": "9s6TZVkeHVmj"
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sJD-i_Vr0vzw"
   },
   "source": [
    "## Before You Begin\n",
    "\n",
    "To run this notebook, you will need to do the following:\n",
    "\n",
    " * [Create a Google Cloud Project](https://developers.google.com/workspace/guides/create-project)\n",
    " * [Enable the AlloyDB API](https://console.cloud.google.com/flows/enableapi?apiid=alloydb.googleapis.com)\n",
    " * [Create a AlloyDB instance](https://cloud.google.com/alloydb/docs/instance-primary-create)\n",
    " * [Create a AlloyDB database](https://cloud.google.com/alloydb/docs/database-create)\n",
    " * [Add an IAM database user to the database](https://cloud.google.com/alloydb/docs/manage-iam-authn) (Optional)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AvefHSmF0vzx"
   },
   "source": [
    "### 🦜🔗 Library Installation\n",
    "The integration lives in its own `langchain-google-alloydb-pg` package, so we need to install it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IFqqkBeO0vzx"
   },
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet langchain-google-alloydb-pg[langgraph] langgraph langchain-google-vertexai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JVdfcv1a0vzy"
   },
   "source": [
    "**Colab only:** Uncomment the following cell to restart the kernel or use the button to restart the kernel. For Vertex AI Workbench you can restart the terminal using the button on top."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tbp91SUI0vzz"
   },
   "outputs": [],
   "source": [
    "# # Automatically restart kernel after installs so that your environment can access the new packages\n",
    "# import IPython\n",
    "\n",
    "# app = IPython.Application.instance()\n",
    "# app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SJF6ZLV30vzz"
   },
   "source": [
    "### 🔐 Authentication\n",
    "Authenticate to Google Cloud as the IAM user logged into this notebook in order to access your Google Cloud Project.\n",
    "\n",
    "* If you are using Colab to run this notebook, use the cell below and continue.\n",
    "* If you are using Vertex AI Workbench, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GmQRkmrb0vzz"
   },
   "outputs": [],
   "source": [
    "from google.colab import auth\n",
    "\n",
    "auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uNDcuuKJ0vzz"
   },
   "source": [
    "### ☁ Set Your Google Cloud Project\n",
    "Set your Google Cloud project so that you can leverage Google Cloud resources within this notebook.\n",
    "\n",
    "If you don't know your project ID, try the following:\n",
    "\n",
    "* Run `gcloud config list`.\n",
    "* Run `gcloud projects list`.\n",
    "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TlMwe3F20vz0",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "ea423aa8-4d5c-4a97-ab97-a474c3a7fa51"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Updated property [core/project].\n"
     ]
    }
   ],
   "source": [
    "# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.\n",
    "\n",
    "PROJECT_ID = \"my-project-id\"  # @param {type:\"string\"}\n",
    "\n",
    "# Set the project id\n",
    "!gcloud config set project {PROJECT_ID}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mLNwdPai0vz0"
   },
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GAq8C9cy0vz0"
   },
   "source": [
    "### Set AlloyDB database values\n",
    "Find your database values, in the [AlloyDB cluster page](https://console.cloud.google.com/alloydb?_ga=2.223735448.2062268965.1707700487-2088871159.1707257687)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iJQzfVk80vz0"
   },
   "outputs": [],
   "source": [
    "# @title Set Your Values Here { display-mode: \"form\" }\n",
    "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
    "CLUSTER = \"my-alloydb-cluster\"  # @param {type: \"string\"}\n",
    "INSTANCE = \"my-alloydb-instance\"  # @param {type: \"string\"}\n",
    "DATABASE = \"my-database\"  # @param {type: \"string\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AXIR0Uaz0vz0"
   },
   "source": [
    "### AlloyDBEngine Connection Pool\n",
    "\n",
    "One of the requirements and arguments to establish AlloyDB as a Checkpoint memory store is a `AlloyDBEngine` object. The `AlloyDBEngine`  configures a connection pool to your AlloyDB database, enabling successful connections from your application and following industry best practices.\n",
    "\n",
    "To create a `AlloyDBEngine` using `AlloyDBEngine.from_instance()` you need to provide only 5 things:\n",
    "\n",
    "1. `project_id` : Project ID of the Google Cloud Project where the AlloyDB instance is located.\n",
    "1. `region` : Region where the AlloyDB instance is located.\n",
    "1. `cluster`: The name of the AlloyDB cluster.\n",
    "1. `instance` : The name of the AlloyDB instance.\n",
    "1. `database` : The name of the database to connect to on the AlloyDB instance.\n",
    "\n",
    "By default, [IAM database authentication](https://cloud.google.com/alloydb/docs/manage-iam-authn) will be used as the method of database authentication. This library uses the IAM principal belonging to the [Application Default Credentials (ADC)](https://cloud.google.com/docs/authentication/application-default-credentials) sourced from the envionment.\n",
    "\n",
    "Optionally, [built-in database authentication](https://cloud.google.com/alloydb/docs/database-users/about) using a username and password to access the AlloyDB database can also be used. Just provide the optional `user` and `password` arguments to `AlloyDBEngine.from_instance()`:\n",
    "\n",
    "* `user` : Database user to use for built-in database authentication and login\n",
    "* `password` : Database password to use for built-in database authentication and login.\n",
    "\n",
    "To connect to your AlloyDB instance from this notebook, you will need to enable public IP on your instance. Alternatively, you can follow [these instructions](https://cloud.google.com/alloydb/docs/connect-external) to connect to an AlloyDB for PostgreSQL instance with Private IP from outside your VPC.\n",
    "Learn more about [specifying IP types](https://github.com/GoogleCloudPlatform/alloydb-python-connector?tab=readme-ov-file#specifying-ip-address-type).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yb2YAOzS0vz0"
   },
   "outputs": [],
   "source": [
    "from langchain_google_alloydb_pg import AlloyDBEngine\n",
    "from google.cloud.alloydb.connector import IPTypes\n",
    "\n",
    "engine = AlloyDBEngine.from_instance(\n",
    "    project_id=PROJECT_ID,\n",
    "    region=REGION,\n",
    "    cluster=CLUSTER,\n",
    "    instance=INSTANCE,\n",
    "    database=DATABASE,\n",
    "    ip_type=IPTypes.PUBLIC,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O3vW0IYF0vz1"
   },
   "source": [
    "### AlloyDBEngine for AlloyDB Omni\n",
    "To create an `AlloyDBEngine` for AlloyDB Omni, you will need a connection url. `AlloyDBEngine.from_engine_args` first creates an async engine and then turns it into an `AlloyDBEngine`. Here is an example connection with the `asyncpg` driver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dqWWKcgj0vz1"
   },
   "outputs": [],
   "source": [
    "# Replace with your own AlloyDB Omni info\n",
    "connstring = f\"postgresql+asyncpg://{OMNI_USER}:{OMNI_PASSWORD}@{OMNI_HOST}:{OMNI_PORT}/{OMNI_DATABASE}\"\n",
    "engine = AlloyDBEngine.from_engine_args(connstring)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eyh2vCQq0vz1"
   },
   "source": [
    "### Initialize a table\n",
    "The `AlloyDBSaver` class requires a database table with a specific schema in order to store the persist LangGraph agents state.\n",
    "The `AlloyDBEngine` engine has a helper method `init_checkpoint_table()` that can be used to create a table with the proper schema for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ehpAbwH50vz1"
   },
   "outputs": [],
   "source": [
    "engine.init_checkpoint_table()  # Use table_name to customise the table name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qdIspizq0vz1"
   },
   "source": [
    "#### Optional Tip: 💡\n",
    "You can also specify a schema name by passing `schema_name` wherever you pass `table_name`. Eg:\n",
    "\n",
    "```python\n",
    "SCHEMA_NAME=\"my_schema\"\n",
    "\n",
    "engine.init_checkpoint_table(\n",
    "    table_name=TABLE_NAME,\n",
    "    schema_name=SCHEMA_NAME     # Default: \"public\"\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a-UERu8q0vz2"
   },
   "source": [
    "### AlloyDBSaver\n",
    "\n",
    "To initialize the `AlloyDBSaver` class you need to provide only 4 things:\n",
    "\n",
    "1. `engine` - An instance of a `AlloyDBEngine` engine.\n",
    "1. `table_name` : Table name that stores the checkpoints (default: \"checkpoints\").\n",
    "1. `schema_name` : The schema name where the table is located (default: \"public\").\n",
    "1. `serde`: Serializer for encoding/decoding checkpoints (default: None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6pr-ZJR-0vz2"
   },
   "outputs": [],
   "source": [
    "from langchain_google_alloydb_pg import AlloyDBSaver\n",
    "\n",
    "checkpointer = AlloyDBSaver.create_sync(\n",
    "    engine,\n",
    "    # table_name=TABLE_NAME,\n",
    "    # schema_name=SCHEMA_NAME,\n",
    "    # serde=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 📔 Example of Checkpointer methods"
   ],
   "metadata": {
    "id": "nx1BkPq-869u"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "write_config = {\"configurable\": {\"thread_id\": \"1\", \"checkpoint_ns\": \"\"}}\n",
    "read_config = {\"configurable\": {\"thread_id\": \"1\"}}\n",
    "\n",
    "checkpoint = {\n",
    "    \"v\": 1,\n",
    "    \"ts\": \"2024-07-31T20:14:19.804150+00:00\",\n",
    "    \"id\": \"1ef4f797-8335-6428-8001-8a1503f9b875\",\n",
    "    \"channel_values\": {\"my_key\": \"meow\", \"node\": \"node\"},\n",
    "    \"channel_versions\": {\"__start__\": 2, \"my_key\": 3, \"start:node\": 3, \"node\": 3},\n",
    "    \"versions_seen\": {\n",
    "        \"__input__\": {},\n",
    "        \"__start__\": {\"__start__\": 1},\n",
    "        \"node\": {\"start:node\": 2},\n",
    "    },\n",
    "    \"pending_sends\": [],\n",
    "}\n",
    "\n",
    "# store checkpoint\n",
    "checkpointer.put(write_config, checkpoint, {}, {})\n",
    "# load checkpoint\n",
    "checkpointer.get(read_config)"
   ],
   "metadata": {
    "id": "zopFUbLz9B4V"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "33xBzm-o0vz3"
   },
   "source": [
    "## 🔗 Adding persistence to the pre-built create react agent\n"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "from typing import Literal\n",
    "\n",
    "from langchain_core.tools import tool\n",
    "from langchain_google_vertexai import ChatVertexAI\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "\n",
    "@tool\n",
    "def get_weather(city: Literal[\"nyc\", \"sf\"]):\n",
    "    \"\"\"Use this to get weather information.\"\"\"\n",
    "    if city == \"nyc\":\n",
    "        return \"It might be cloudy in nyc\"\n",
    "    elif city == \"sf\":\n",
    "        return \"It's always sunny in sf\"\n",
    "    else:\n",
    "        raise AssertionError(\"Unknown city\")\n",
    "\n",
    "\n",
    "tools = [get_weather]\n",
    "model = ChatVertexAI(project=PROJECT_ID, model_name=\"gemini-2.0-flash-exp\")\n",
    "\n",
    "graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n",
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n",
    "print(res)"
   ],
   "metadata": {
    "id": "FFqcTadHRs1S"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# Example of resulting checkpoint config\n",
    "checkpoint = checkpointer.get(config)"
   ],
   "metadata": {
    "id": "M-8JEd2qS4QQ"
   },
   "execution_count": null,
   "outputs": []
  }
 ],
 "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.11.5"
  },
  "colab": {
   "provenance": [],
   "collapsed_sections": [
    "a-UERu8q0vz2"
   ]
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}