docs/graph_store.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7VBkjcqNNxEd"
   },
   "source": [
    "# Google Spanner\n",
    "\n",
    "> [Spanner](https://cloud.google.com/spanner) is a highly scalable database that combines unlimited scalability with relational semantics, such as secondary indexes, strong consistency, schemas, and SQL providing 99.999% availability in one easy solution.\n",
    "\n",
    "This notebook goes over how to use `Spanner` for GraphRAG with `SpannerGraphStore` class.\n",
    "\n",
    "Learn more about the package on [GitHub](https://github.com/googleapis/langchain-google-spanner-python/).\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/googleapis/langchain-google-spanner-python/blob/main/docs/graph_store.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HEAGYTPgNydh"
   },
   "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 Cloud Spanner API](https://console.cloud.google.com/flows/enableapi?apiid=spanner.googleapis.com)\n",
    " * [Create a Spanner instance](https://cloud.google.com/spanner/docs/create-manage-instances)\n",
    " * [Create a Spanner database](https://cloud.google.com/spanner/docs/create-manage-databases)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cboPIg-yOcxS"
   },
   "source": [
    "### 🦜🔗 Library Installation\n",
    "The integration lives in its own `langchain-google-spanner` package, so we need to install it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "AOWh6QKYVdDp"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --upgrade --quiet langchain-google-spanner langchain-google-vertexai langchain-experimental json-repair wikipedia pyvis youtube-transcript-api"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M7MqpDhkOiP-"
   },
   "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": "xzgVZv0POj17",
    "outputId": "066ba1ae-c89e-427e-8a98-7f76f90155d0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'status': 'ok', 'restart': True}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": "zfIhwIryOls1"
   },
   "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": "EWOkHI7XOna2"
   },
   "outputs": [],
   "source": [
    "from google.colab import auth\n",
    "\n",
    "auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6xHXneICOpsB"
   },
   "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": "hF0481BGOsS8"
   },
   "outputs": [],
   "source": [
    "# @markdown Please fill in the value below with your Google Cloud project ID and then run the cell.\n",
    "\n",
    "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
    "\n",
    "# Set the project id\n",
    "!gcloud config set project {PROJECT_ID}\n",
    "%env GOOGLE_CLOUD_PROJECT={PROJECT_ID}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4TiC0RbhOwUu"
   },
   "source": [
    "### 💡 API Enablement\n",
    "The `langchain-google-spanner` package requires that you [enable the Spanner API](https://console.cloud.google.com/flows/enableapi?apiid=spanner.googleapis.com) in your Google Cloud Project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9f3fJd5eOyRr"
   },
   "outputs": [],
   "source": [
    "# enable Spanner API\n",
    "!gcloud services enable spanner.googleapis.com"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "k5pxMMiMOzt7"
   },
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G7-Pe2ADQlNJ"
   },
   "source": [
    "### Prepare graph documents\n",
    "\n",
    "Prepare graph documents from various sources to be added to Spanner Graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0rdoT01aRiNI",
    "outputId": "c63f6325-3e64-481e-e3d8-a3cb73be2d95"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fetching document from Wikipedia: `Google`\n",
      "Fetching document from YouTube: `https://www.youtube.com/watch?v=WsEQjeZoEng`\n"
     ]
    }
   ],
   "source": [
    "# Get graph documents from Wikipedia & YouTube\n",
    "from langchain_community.document_loaders import WikipediaLoader, YoutubeLoader\n",
    "from langchain_experimental.graph_transformers import LLMGraphTransformer\n",
    "from langchain_google_vertexai.llms import VertexAI\n",
    "\n",
    "llm = VertexAI(model_name=\"gemini-2.0-flash-001\")\n",
    "llm_transformer = LLMGraphTransformer(\n",
    "    llm=llm, allowed_nodes=[\"Person\", \"Company\", \"Product\"]\n",
    ")\n",
    "\n",
    "\n",
    "def get_document_from_wiki(query):\n",
    "    print(\"Fetching document from Wikipedia: `%s`\" % query)\n",
    "    loader = WikipediaLoader(query=query, load_max_docs=1)\n",
    "    return loader.load()\n",
    "\n",
    "\n",
    "def get_document_from_youtube(link):\n",
    "    print(\"Fetching document from YouTube: `%s`\" % link)\n",
    "    loader = YoutubeLoader.from_youtube_url(link)\n",
    "    return loader.load()\n",
    "\n",
    "\n",
    "document_batch = list(map(get_document_from_wiki, [\"Google\"]))\n",
    "document_batch += list(\n",
    "    map(get_document_from_youtube, [\"https://www.youtube.com/watch?v=WsEQjeZoEng\"])\n",
    ")\n",
    "graph_documents = llm_transformer.convert_to_graph_documents(\n",
    "    [document for documents in document_batch for document in documents]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OylyNyv-ZsT2",
    "outputId": "875bf1c3-62e5-4a0b-d064-1900ecc8d405"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Google LLC ( , GOO-gəl) is an American multinational corporation and technology company focusing on \n",
      "[Node(id='Google LLC', type='Company', properties={}), Node(id='Larry Page', type='Person', properties={}), Node(id='Sergey Brin', type='Person', properties={})]\n",
      "[Relationship(source=Node(id='Google LLC', type='Company', properties={}), target=Node(id='Larry Page', type='Person', properties={}), type='FOUNDED_BY', properties={}), Relationship(source=Node(id='Google LLC', type='Company', properties={}), target=Node(id='Sergey Brin', type='Person', properties={}), type='FOUNDED_BY', properties={})]\n",
      "\n",
      "[MUSIC PLAYING] [APPLAUSE] SUNDAR PICHAI: At Google, we\n",
      "are fully in our Gemini era. Today, all of o\n",
      "[Node(id='Google', type='Company', properties={}), Node(id='Gemini', type='Product', properties={}), Node(id='summarize emails', type='Feature', properties={})]\n",
      "[Relationship(source=Node(id='Google', type='Company', properties={}), target=Node(id='Gemini', type='Product', properties={}), type='HAS_PRODUCT', properties={}), Relationship(source=Node(id='Gemini', type='Product', properties={}), target=Node(id='summarize emails', type='Feature', properties={}), type='HAS_FEATURE', properties={})]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for doc in graph_documents:\n",
    "    print(doc.source.page_content[:100])\n",
    "    print(doc.nodes)\n",
    "    print(doc.relationships)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mtDbLU5sO2iA"
   },
   "source": [
    "### Set Spanner database values\n",
    "Find your database values, in the [Spanner Instances page](https://console.cloud.google.com/spanner?_ga=2.223735448.2062268965.1707700487-2088871159.1707257687).\n",
    "\n",
    "NOTE:\n",
    "- The database identified by INSTANCE and DATABASE must be created beforehand.\n",
    "- The graph does NOT need to be created beforehand.\n",
    "  \n",
    "  If such a graph already exists, graphs will be built atop; otherwise a new graph is automatically created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "C-I8VTIcO442"
   },
   "outputs": [],
   "source": [
    "# @title Set Your Values Here { display-mode: \"form\" }\n",
    "INSTANCE = \"\"  # @param {type: \"string\"}\n",
    "DATABASE = \"\"  # @param {type: \"string\"}\n",
    "GRAPH_NAME = \"\"  # @param {type: \"string\"}\n",
    "USE_FLEXIBLE_SCHEMA = False  # @param {type: \"boolean\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kpAv-tpcO_iL"
   },
   "source": [
    "### SpannerGraphStore\n",
    "\n",
    "To initialize the `SpannerGraphStore` class you need to provide 3 required arguments and other arguments are optional and only need to pass if it's different from default ones\n",
    "\n",
    "1.   a Spanner instance id;\n",
    "2.   a Spanner database id belongs to the above instance id;\n",
    "3.   a Spanner graph name used to create a graph in the above database.\n",
    "\n",
    "#### SpannerGraphStore with flexible schema\n",
    "\n",
    "By default, SpannerGraphStore creates an underlying table for each type of nodes and edges.\n",
    "This will create many underlying tables when your graph consists of many different types of nodes and edges.\n",
    "\n",
    "SpannerGraphStore provides a flexible schema mode that stores all your nodes in a single node table and edges in a single edge table.\n",
    "\n",
    "To use it, set USE_FLEXIBLE_SCHEMA to True."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "id": "u589YapWQFb8"
   },
   "outputs": [],
   "source": [
    "from langchain_google_spanner import SpannerGraphStore\n",
    "\n",
    "graph_store = SpannerGraphStore(\n",
    "    instance_id=INSTANCE,\n",
    "    database_id=DATABASE,\n",
    "    graph_name=GRAPH_NAME,\n",
    "    use_flexible_schema=USE_FLEXIBLE_SCHEMA,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lMXvOpRbZdau",
    "outputId": "3c6ca58c-0e6a-4b7f-a204-eab08adb2761"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Clean up existing data...\n",
      "Waiting for DDL operations to complete...\n",
      "Waiting for DDL operations to complete...\n",
      "Waiting for DDL operations to complete...\n",
      "Add new data...\n",
      "Waiting for DDL operations to complete...\n",
      "Insert nodes of type `Company`...\n",
      "Insert nodes of type `Person`...\n",
      "Insert edges of type `Company_FOUNDED_BY_Person`...\n",
      "No schema change required...\n",
      "Insert nodes of type `Company`...\n",
      "Insert nodes of type `Product`...\n",
      "Insert nodes of type `Feature`...\n",
      "Insert edges of type `Company_HAS_PRODUCT_Product`...\n",
      "Insert edges of type `Product_HAS_FEATURE_Feature`...\n"
     ]
    }
   ],
   "source": [
    "print(\"Clean up existing data...\")\n",
    "graph_store.cleanup()\n",
    "\n",
    "print(\"Add new data...\")\n",
    "for graph_document in graph_documents:\n",
    "    graph_store.add_graph_documents([graph_document])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v5b572QbSi4t"
   },
   "source": [
    "#### Query the graph\n",
    "To traverse the graph in the graph store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9YpR4cauSsfG",
    "outputId": "387bb32b-c28b-4be7-efe8-7797c5914d1c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'google_related_nodes': ['Larry Page', 'Gemini', 'Google', 'summarize emails', 'Sergey Brin', 'Google LLC']}]\n"
     ]
    }
   ],
   "source": [
    "sample_query = \"\"\"\n",
    "  GRAPH {}\n",
    "  MATCH (n WHERE REGEXP_CONTAINS(n.id, 'Google')) -[e]-{{1, 2}} (m)\n",
    "  RETURN ARRAY_AGG(DISTINCT m.id) AS google_related_nodes\n",
    "\"\"\".format(\n",
    "    GRAPH_NAME\n",
    ")\n",
    "\n",
    "print(graph_store.query(sample_query))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xau7BXvb_6si"
   },
   "source": [
    "#### Visualize the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 551
    },
    "id": "pQuQzOng_9GK",
    "outputId": "2d791b28-f51a-4c20-de3b-6293fd84070d"
   },
   "outputs": [],
   "source": [
    "from pyvis.network import Network\n",
    "from IPython.core.display import display, HTML\n",
    "\n",
    "net = Network(\n",
    "    notebook=True,\n",
    "    cdn_resources=\"remote\",\n",
    "    bgcolor=\"#222222\",\n",
    "    font_color=\"white\",\n",
    "    height=\"500px\",\n",
    "    width=\"50%\",\n",
    "    directed=True,\n",
    ")\n",
    "\n",
    "node_query = \"\"\"\n",
    "  GRAPH {}\n",
    "  MATCH (n)\n",
    "  RETURN n.id\n",
    "\"\"\".format(\n",
    "    GRAPH_NAME\n",
    ")\n",
    "\n",
    "edge_query = \"\"\"\n",
    "  GRAPH {}\n",
    "  MATCH -[e]->\n",
    "  RETURN e.id AS src_id, e.target_id AS dst_id, labels(e)[0] AS label\n",
    "\"\"\".format(\n",
    "    GRAPH_NAME\n",
    ")\n",
    "\n",
    "nodes = graph_store.query(node_query)\n",
    "edges = graph_store.query(edge_query)\n",
    "\n",
    "net.add_nodes([node[\"id\"] for node in nodes])\n",
    "for edge in edges:\n",
    "    net.add_edge(edge[\"src_id\"], edge[\"dst_id\"], title=edge[\"label\"])\n",
    "\n",
    "net.toggle_physics(True)\n",
    "net.show(\"graph.html\")\n",
    "display(HTML(\"graph.html\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pM7TmfI0TEFy"
   },
   "source": [
    "#### Clean up the graph\n",
    "\n",
    "> USE IT WITH CAUTION!\n",
    "\n",
    "Clean up all the nodes/edges in your graph and remove your graph definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UQWq4-sITOgl",
    "outputId": "5ed5ebf2-a34b-400e-a171-71b2e863e99c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Waiting for DDL operations to complete...\n",
      "Waiting for DDL operations to complete...\n",
      "Waiting for DDL operations to complete...\n"
     ]
    }
   ],
   "source": [
    "graph_store.cleanup()"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}