docs/graph_qa_chain.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "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": [],
   "toc_visible": true
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "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` and `SpannerGraphQAChain` 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_qa_chain.ipynb)"
   ],
   "metadata": {
    "id": "7VBkjcqNNxEd"
   }
  },
  {
   "cell_type": "markdown",
   "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)"
   ],
   "metadata": {
    "id": "HEAGYTPgNydh"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 🦜🔗 Library Installation\n",
    "The integration lives in its own `langchain-google-spanner` package, so we need to install it."
   ],
   "metadata": {
    "id": "cboPIg-yOcxS"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "%pip install --upgrade --quiet langchain-google-spanner langchain-google-vertexai langchain-experimental json-repair pyvis"
   ],
   "metadata": {
    "id": "AOWh6QKYVdDp"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "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."
   ],
   "metadata": {
    "id": "M7MqpDhkOiP-"
   }
  },
  {
   "cell_type": "code",
   "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)"
   ],
   "metadata": {
    "id": "xzgVZv0POj17"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "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)."
   ],
   "metadata": {
    "id": "zfIhwIryOls1"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "from google.colab import auth\n",
    "\n",
    "auth.authenticate_user()"
   ],
   "metadata": {
    "id": "EWOkHI7XOna2"
   },
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "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)."
   ],
   "metadata": {
    "id": "6xHXneICOpsB"
   }
  },
  {
   "cell_type": "code",
   "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}\n",
    "%env GOOGLE_CLOUD_PROJECT={PROJECT_ID}"
   ],
   "metadata": {
    "id": "hF0481BGOsS8"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "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."
   ],
   "metadata": {
    "id": "4TiC0RbhOwUu"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# enable Spanner API\n",
    "!gcloud services enable spanner.googleapis.com"
   ],
   "metadata": {
    "id": "9f3fJd5eOyRr"
   },
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "You must also and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)."
   ],
   "metadata": {
    "id": "bT_S-jaEOW4P"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Basic Usage"
   ],
   "metadata": {
    "id": "k5pxMMiMOzt7"
   }
  },
  {
   "cell_type": "markdown",
   "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)."
   ],
   "metadata": {
    "id": "mtDbLU5sO2iA"
   }
  },
  {
   "cell_type": "code",
   "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\"}"
   ],
   "metadata": {
    "id": "C-I8VTIcO442"
   },
   "execution_count": 15,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "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."
   ],
   "metadata": {
    "id": "kpAv-tpcO_iL"
   }
  },
  {
   "cell_type": "code",
   "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",
    ")"
   ],
   "metadata": {
    "id": "u589YapWQFb8"
   },
   "execution_count": 16,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### Add Graph Documents to Spanner Graph"
   ],
   "metadata": {
    "id": "G7-Pe2ADQlNJ"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# @title Extract Nodes and Edges from text snippets\n",
    "from langchain_core.documents import Document\n",
    "from langchain_experimental.graph_transformers import LLMGraphTransformer\n",
    "from langchain_google_vertexai import ChatVertexAI\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_snippets = [\n",
    "    # Text snippet for students graduting from Veritas University, Computer Science Dept 2017\n",
    "    \"\"\"\n",
    "This was the graduation ceremony of 2017. A wave of jubilant graduates poured out of the\n",
    "grand halls of Veritas University, their laughter echoing across the quad. Among them were\n",
    "a cohort of exceptional students from the Computer Science department, a group that had\n",
    "become known for their collaborative spirit and innovative ideas.\n",
    "Leading the pack was Emily Davis, a coding whiz with a passion for cybersecurity, already\n",
    "fielding offers from top tech firms. Beside her walked James Rodriguez, a quiet but\n",
    "brilliant mind fascinated by artificial intelligence, dreaming of building machines that\n",
    "could understand human emotions.  Trailing slightly behind, deep in conversation, were\n",
    "Sarah Chen and Michael Patel, both aspiring game developers, eager to bring their creative\n",
    "visions to life.  And then there was  Aisha Khan, a social justice advocate who planned to\n",
    "use her coding skills to address inequality through technology.\n",
    "As they celebrated their achievements, these Veritas University Computer Science graduates\n",
    "were ready to embark on diverse paths, each carrying the potential to shape the future of\n",
    "technology in their own unique way.\n",
    "\"\"\",\n",
    "    # Text snippet for students graduting from Oakhaven University, Computer Science Dept 2016\n",
    "    \"\"\"\n",
    "The year was 2016, and a palpable buzz filled the air as the graduating class of Oakhaven\n",
    "university from Computer science and Engineering department emerged from the Beckman\n",
    "Auditorium. Among them was a group of exceptional students, renowned for their\n",
    "intellectual curiosity and groundbreaking research.\n",
    "At the forefront was Alice Johnson, a gifted programmer with a fascination for quantum\n",
    "computing, already collaborating with leading researchers in the field.  Beside her\n",
    "strode David Kim, a brilliant theorist captivated by the intricacies of cryptography,\n",
    "eager to contribute to the development of secure communication systems.  Engaged in an\n",
    "animated discussion were Maria Rodriguez and Robert Lee, both passionate about robotics\n",
    "and determined to push the boundaries of artificial intelligence.  And then there was\n",
    "Chloe Brown, a visionary with a deep interest in bioinformatics, driven to unlock the\n",
    "secrets of the human genome through computational analysis.\n",
    "As they celebrated their accomplishments, these graduates, armed with their exceptional\n",
    "skills and unwavering determination, were poised to make significant contributions to the world of computing and beyond.\n",
    "\"\"\",\n",
    "    # Text snippet mentions the company Emily Davis founded.\n",
    "    # The snippet doesn't mention that she is an alumni of Veritas University\n",
    "    \"\"\"\n",
    "Emily Davis, a name synonymous with cybersecurity innovation, turned that passion into a\n",
    "thriving business.  In the year 2022, Davis founded Ironclad Security, a company that's\n",
    "rapidly changing the landscape of cybersecurity solutions.\n",
    "\"\"\",\n",
    "    # Text snippet mentions the company Alice Johnson founded.\n",
    "    # The snippet doesn't mention that she is an alumni of Oakhaven University.\n",
    "    \"\"\"\n",
    "Alice Johnson had a vision that extended far beyond the classroom. Driven by an insatiable\n",
    "curiosity about the potential of quantum mechanics, she founded Entangled Solutions, a\n",
    "company poised to revolutionize industries through the power of quantum technology.\n",
    "Entangled Solutions distinguishes itself by focusing on practical applications of quantum\n",
    "computing.\n",
    "\"\"\",\n",
    "]\n",
    "\n",
    "# Create splits for documents\n",
    "documents = [Document(page_content=t) for t in text_snippets]\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n",
    "splits = text_splitter.split_documents(documents)\n",
    "llm = ChatVertexAI(model=\"gemini-1.5-flash\", temperature=0)\n",
    "llm_transformer = LLMGraphTransformer(\n",
    "    llm=llm,\n",
    "    allowed_nodes=[\"College\", \"Deparatment\", \"Person\", \"Year\", \"Company\"],\n",
    "    allowed_relationships=[\n",
    "        \"AlumniOf\",\n",
    "        \"StudiedInDepartment\",\n",
    "        \"PartOf\",\n",
    "        \"GraduatedInYear\",\n",
    "        \"Founded\",\n",
    "    ],\n",
    "    node_properties=[\n",
    "        \"description\",\n",
    "    ],\n",
    ")\n",
    "graph_documents = llm_transformer.convert_to_graph_documents(splits)"
   ],
   "metadata": {
    "id": "fP7XNu3aPl5c"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# @title Print extracted nodes and edges\n",
    "for doc in graph_documents:\n",
    "    print(doc.source.page_content[:100])\n",
    "    print(doc.nodes)\n",
    "    print(doc.relationships)\n",
    "    print()"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OylyNyv-ZsT2",
    "outputId": "e4253d98-ad63-4ea8-a5f1-0e3dac8f6632"
   },
   "execution_count": 13,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "This was the graduation ceremony of 2017. A wave of jubilant graduates poured out of the\n",
      "grand halls\n",
      "[Node(id='Veritas University', type='College', properties={'description': 'grand halls'}), Node(id='Computer Science', type='Deparatment', properties={}), Node(id='2017', type='Year', properties={}), Node(id='Emily Davis', type='Person', properties={'description': 'coding whiz with a passion for cybersecurity'}), Node(id='James Rodriguez', type='Person', properties={'description': 'quiet but brilliant mind fascinated by artificial intelligence'}), Node(id='Sarah Chen', type='Person', properties={'description': 'aspiring game developers'}), Node(id='Michael Patel', type='Person', properties={'description': 'aspiring game developers'}), Node(id='Aisha Khan', type='Person', properties={'description': 'social justice advocate'})]\n",
      "[Relationship(source=Node(id='Emily Davis', type='Person', properties={}), target=Node(id='Veritas University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='James Rodriguez', type='Person', properties={}), target=Node(id='Veritas University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Sarah Chen', type='Person', properties={}), target=Node(id='Veritas University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Michael Patel', type='Person', properties={}), target=Node(id='Veritas University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Aisha Khan', type='Person', properties={}), target=Node(id='Veritas University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Emily Davis', type='Person', properties={}), target=Node(id='Computer Science', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='James Rodriguez', type='Person', properties={}), target=Node(id='Computer Science', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Sarah Chen', type='Person', properties={}), target=Node(id='Computer Science', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Michael Patel', type='Person', properties={}), target=Node(id='Computer Science', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Aisha Khan', type='Person', properties={}), target=Node(id='Computer Science', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Emily Davis', type='Person', properties={}), target=Node(id='2017', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='James Rodriguez', type='Person', properties={}), target=Node(id='2017', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Sarah Chen', type='Person', properties={}), target=Node(id='2017', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Michael Patel', type='Person', properties={}), target=Node(id='2017', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Aisha Khan', type='Person', properties={}), target=Node(id='2017', type='Year', properties={}), type='GRADUATEDINYEAR', properties={})]\n",
      "\n",
      "visions to life.  And then there was  Aisha Khan, a social justice advocate who planned to\n",
      "use her c\n",
      "[Node(id='Veritas University', type='College', properties={}), Node(id='Computer Science', type='Deparatment', properties={}), Node(id='Aisha Khan', type='Person', properties={'description': 'social justice advocate'})]\n",
      "[Relationship(source=Node(id='Aisha Khan', type='Person', properties={}), target=Node(id='Veritas University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Aisha Khan', type='Person', properties={}), target=Node(id='Computer Science', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={})]\n",
      "\n",
      "The year was 2016, and a palpable buzz filled the air as the graduating class of Oakhaven\n",
      "university\n",
      "[Node(id='Oakhaven University', type='College', properties={'description': 'Oakhaven university'}), Node(id='Computer Science And Engineering', type='Deparatment', properties={'description': 'Computer science and Engineering'}), Node(id='2016', type='Year', properties={'description': '2016'}), Node(id='Alice Johnson', type='Person', properties={'description': 'a gifted programmer with a fascination for quantum computing, already collaborating with leading researchers in the field'}), Node(id='David Kim', type='Person', properties={'description': 'a brilliant theorist captivated by the intricacies of cryptography, eager to contribute to the development of secure communication systems'}), Node(id='Maria Rodriguez', type='Person', properties={'description': 'passionate about robotics and determined to push the boundaries of artificial intelligence'}), Node(id='Robert Lee', type='Person', properties={'description': 'passionate about robotics and determined to push the boundaries of artificial intelligence'}), Node(id='Chloe Brown', type='Person', properties={'description': 'a visionary with a deep interest in bioinformatics, driven to unlock the secrets of the human genome through computational analysis'}), Node(id='Beckman Auditorium', type='Deparatment', properties={'description': 'Beckman Auditorium'})]\n",
      "[Relationship(source=Node(id='Alice Johnson', type='Person', properties={}), target=Node(id='Oakhaven University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='David Kim', type='Person', properties={}), target=Node(id='Oakhaven University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Maria Rodriguez', type='Person', properties={}), target=Node(id='Oakhaven University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Robert Lee', type='Person', properties={}), target=Node(id='Oakhaven University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Chloe Brown', type='Person', properties={}), target=Node(id='Oakhaven University', type='College', properties={}), type='ALUMNIOF', properties={}), Relationship(source=Node(id='Computer Science And Engineering', type='Deparatment', properties={}), target=Node(id='Oakhaven University', type='College', properties={}), type='PARTOF', properties={}), Relationship(source=Node(id='Alice Johnson', type='Person', properties={}), target=Node(id='Computer Science And Engineering', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='David Kim', type='Person', properties={}), target=Node(id='Computer Science And Engineering', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Maria Rodriguez', type='Person', properties={}), target=Node(id='Computer Science And Engineering', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Robert Lee', type='Person', properties={}), target=Node(id='Computer Science And Engineering', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Chloe Brown', type='Person', properties={}), target=Node(id='Computer Science And Engineering', type='Deparatment', properties={}), type='STUDIEDINDEPARTMENT', properties={}), Relationship(source=Node(id='Oakhaven University', type='College', properties={}), target=Node(id='2016', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Alice Johnson', type='Person', properties={}), target=Node(id='2016', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='David Kim', type='Person', properties={}), target=Node(id='2016', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Maria Rodriguez', type='Person', properties={}), target=Node(id='2016', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Robert Lee', type='Person', properties={}), target=Node(id='2016', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Chloe Brown', type='Person', properties={}), target=Node(id='2016', type='Year', properties={}), type='GRADUATEDINYEAR', properties={}), Relationship(source=Node(id='Oakhaven University', type='College', properties={}), target=Node(id='Beckman Auditorium', type='Deparatment', properties={}), type='PARTOF', properties={})]\n",
      "\n",
      "Chloe Brown, a visionary with a deep interest in bioinformatics, driven to unlock the\n",
      "secrets of the\n",
      "[Node(id='Chloe Brown', type='Person', properties={'description': 'a visionary with a deep interest in bioinformatics, driven to unlock the secrets of the human genome through computational analysis'})]\n",
      "[]\n",
      "\n",
      "Emily Davis, a name synonymous with cybersecurity innovation, turned that passion into a\n",
      "thriving bu\n",
      "[Node(id='Emily Davis', type='Person', properties={'description': 'a name synonymous with cybersecurity innovation'}), Node(id='Ironclad Security', type='Company', properties={'description': \"a company that's rapidly changing the landscape of cybersecurity solutions\"}), Node(id='2022', type='Year', properties={})]\n",
      "[Relationship(source=Node(id='Emily Davis', type='Person', properties={}), target=Node(id='Ironclad Security', type='Company', properties={}), type='FOUNDED', properties={}), Relationship(source=Node(id='Emily Davis', type='Person', properties={}), target=Node(id='2022', type='Year', properties={}), type='FOUNDED', properties={})]\n",
      "\n",
      "Alice Johnson had a vision that extended far beyond the classroom. Driven by an insatiable\n",
      "curiosity\n",
      "[Node(id='Alice Johnson', type='Person', properties={'description': 'Driven by an insatiable curiosity about the potential of quantum mechanics, she founded Entangled Solutions, a company poised to revolutionize industries through the power of quantum technology.'}), Node(id='Entangled Solutions', type='Company', properties={'description': 'Entangled Solutions distinguishes itself by focusing on practical applications of quantum computing.'})]\n",
      "[Relationship(source=Node(id='Alice Johnson', type='Person', properties={}), target=Node(id='Entangled Solutions', type='Company', properties={}), type='FOUNDED', properties={})]\n",
      "\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# @title Load the graph to Spanner Graph database\n",
    "# Uncomment the line below, if you want to cleanup from\n",
    "# previous iterations.\n",
    "# BeWARE - THIS COULD REMOVE DATA FROM YOUR DATABASE !!!\n",
    "# graph_store.cleanup()\n",
    "\n",
    "\n",
    "for graph_document in graph_documents:\n",
    "    graph_store.add_graph_documents([graph_document])"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lMXvOpRbZdau",
    "outputId": "26647456-2316-46e3-de43-cfc9845a1050"
   },
   "execution_count": 18,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Waiting for DDL operations to complete...\n",
      "Insert nodes of type `College`...\n",
      "Insert nodes of type `Deparatment`...\n",
      "Insert nodes of type `Year`...\n",
      "Insert nodes of type `Person`...\n",
      "Insert edges of type `Person_ALUMNIOF_College`...\n",
      "Insert edges of type `Person_STUDIEDINDEPARTMENT_Deparatment`...\n",
      "Insert edges of type `Person_GRADUATEDINYEAR_Year`...\n",
      "No schema change required...\n",
      "Insert nodes of type `College`...\n",
      "Insert nodes of type `Deparatment`...\n",
      "Insert nodes of type `Person`...\n",
      "Insert edges of type `Person_ALUMNIOF_College`...\n",
      "Insert edges of type `Person_STUDIEDINDEPARTMENT_Deparatment`...\n",
      "Waiting for DDL operations to complete...\n",
      "Insert nodes of type `College`...\n",
      "Insert nodes of type `Deparatment`...\n",
      "Insert nodes of type `Year`...\n",
      "Insert nodes of type `Person`...\n",
      "Insert edges of type `Person_ALUMNIOF_College`...\n",
      "Insert edges of type `Deparatment_PARTOF_College`...\n",
      "Insert edges of type `Person_STUDIEDINDEPARTMENT_Deparatment`...\n",
      "Insert edges of type `College_GRADUATEDINYEAR_Year`...\n",
      "Insert edges of type `Person_GRADUATEDINYEAR_Year`...\n",
      "Insert edges of type `College_PARTOF_Deparatment`...\n",
      "No schema change required...\n",
      "Insert nodes of type `Person`...\n",
      "Waiting for DDL operations to complete...\n",
      "Insert nodes of type `Person`...\n",
      "Insert nodes of type `Company`...\n",
      "Insert nodes of type `Year`...\n",
      "Insert edges of type `Person_FOUNDED_Company`...\n",
      "Insert edges of type `Person_FOUNDED_Year`...\n",
      "No schema change required...\n",
      "Insert nodes of type `Person`...\n",
      "Insert nodes of type `Company`...\n",
      "Insert edges of type `Person_FOUNDED_Company`...\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Initialize the Spanner Graph QA Chain\n",
    "The Spanner Graph QA Chain takes two parameters, a SpannerGraphStore object and a language model."
   ],
   "metadata": {
    "id": "qlKwtdGN7kaT"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "from google.cloud import spanner\n",
    "from langchain_google_spanner import SpannerGraphQAChain\n",
    "from langchain_google_vertexai import ChatVertexAI\n",
    "from IPython.core.display import HTML\n",
    "\n",
    "# Initialize llm object\n",
    "llm = ChatVertexAI(model=\"gemini-1.5-flash-002\", temperature=0)\n",
    "\n",
    "# Initialize GraphQAChain\n",
    "chain = SpannerGraphQAChain.from_llm(\n",
    "    llm,\n",
    "    graph=graph_store,\n",
    "    allow_dangerous_requests=True,\n",
    "    verbose=True,\n",
    "    return_intermediate_steps=True,\n",
    ")"
   ],
   "metadata": {
    "id": "7yKDAD9s7t7O"
   },
   "execution_count": 30,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "# @title Run Spanner Graph QA Chain 1\n",
    "question = \"Who are the alumni of the college id Veritas University ?\"  # @param {type:\"string\"}\n",
    "response = chain.invoke(\"query=\" + question)\n",
    "response[\"result\"]"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "id": "ukKi9wtH_bF1",
    "outputId": "61b66dcb-54cf-4620-a097-b4f0d732d1e3"
   },
   "execution_count": 33,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SpannerGraphQAChain chain...\u001b[0m\n",
      "Executing gql:\n",
      "\u001b[32;1m\u001b[1;3mGRAPH graph_demo_2\n",
      "MATCH (p:Person)-[:ALUMNIOF]->(c:College {id: \"Veritas University\"})\n",
      "RETURN p.id AS person_id, c.id AS college_id\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'person_id': 'Aisha Khan', 'college_id': 'Veritas University'}, {'person_id': 'Emily Davis', 'college_id': 'Veritas University'}, {'person_id': 'James Rodriguez', 'college_id': 'Veritas University'}, {'person_id': 'Michael Patel', 'college_id': 'Veritas University'}, {'person_id': 'Sarah Chen', 'college_id': 'Veritas University'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'Aisha Khan, Emily Davis, James Rodriguez, Michael Patel, and Sarah Chen are alumni of Veritas University.\\n'"
      ],
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      }
     },
     "metadata": {},
     "execution_count": 33
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# @title Run Spanner Graph QA Chain 2\n",
    "question = \"List the companies, their founders and the college they attended.\"  # @param {type:\"string\"}\n",
    "response = chain.invoke(\"query=\" + question)\n",
    "response[\"result\"]"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 281
    },
    "outputId": "e47d9f63-6769-49bc-b3a3-412c10de5c8a",
    "id": "lcBc4tG__7Rm"
   },
   "execution_count": 34,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SpannerGraphQAChain chain...\u001b[0m\n",
      "Executing gql:\n",
      "\u001b[32;1m\u001b[1;3mGRAPH graph_demo_2\n",
      "MATCH (p:Person)-[:FOUNDED]->(c:Company), (p)-[:ALUMNIOF]->(cl:College)\n",
      "RETURN c.id AS company_id, c.description AS company_description, p.id AS founder_id, p.description AS founder_description, cl.id AS college_id, cl.description AS college_description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'company_id': 'Entangled Solutions', 'company_description': 'Entangled Solutions distinguishes itself by focusing on practical applications of quantum computing.', 'founder_id': 'Alice Johnson', 'founder_description': 'Driven by an insatiable curiosity about the potential of quantum mechanics, she founded Entangled Solutions, a company poised to revolutionize industries through the power of quantum technology.', 'college_id': 'Oakhaven University', 'college_description': 'Oakhaven university'}, {'company_id': 'Ironclad Security', 'company_description': \"a company that's rapidly changing the landscape of cybersecurity solutions\", 'founder_id': 'Emily Davis', 'founder_description': 'a name synonymous with cybersecurity innovation', 'college_id': 'Veritas University', 'college_description': 'grand halls'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'Entangled Solutions, founded by Alice Johnson who attended Oakhaven University, focuses on practical applications of quantum computing.  Ironclad Security, founded by Emily Davis who attended Veritas University, is rapidly changing the landscape of cybersecurity solutions.\\n'"
      ],
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      }
     },
     "metadata": {},
     "execution_count": 34
    }
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# @title Run Spanner Graph QA Chain 3\n",
    "question = \"Which companies were founded by alumni of college id Veritas University ? Who were the founders ?\"  # @param {type:\"string\"}\n",
    "response = chain.invoke(\"query=\" + question)\n",
    "response[\"result\"]"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 264
    },
    "outputId": "cb40179e-bcec-4399-df9d-a114e02b33f9",
    "id": "e6djmq1NAGOM"
   },
   "execution_count": 35,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SpannerGraphQAChain chain...\u001b[0m\n",
      "Executing gql:\n",
      "\u001b[32;1m\u001b[1;3mGRAPH graph_demo_2\n",
      "MATCH (c:College {id: \"Veritas University\"})<-[:ALUMNIOF]-(p:Person)-[:FOUNDED]->(co:Company)\n",
      "RETURN co.id AS company_id, co.description AS company_description, p.id AS founder_id, p.description AS founder_description\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m[{'company_id': 'Ironclad Security', 'company_description': \"a company that's rapidly changing the landscape of cybersecurity solutions\", 'founder_id': 'Emily Davis', 'founder_description': 'a name synonymous with cybersecurity innovation'}]\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "\"Ironclad Security, a company that's rapidly changing the landscape of cybersecurity solutions, was founded by Emily Davis, a name synonymous with cybersecurity innovation.\\n\""
      ],
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      }
     },
     "metadata": {},
     "execution_count": 35
    }
   ]
  },
  {
   "cell_type": "markdown",
   "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."
   ],
   "metadata": {
    "id": "pM7TmfI0TEFy"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "graph_store.cleanup()"
   ],
   "metadata": {
    "id": "UQWq4-sITOgl"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}