run-llama--llama_index/docs/examples/index_structs/knowledge_graph/NeptuneDatabaseKGIndexDemo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "27bc87b7",
   "metadata": {},
   "source": [
    "# Amazon Neptune Graph Store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78b60432",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install boto3\n",
    "%pip install llama-index-llms-bedrock\n",
    "%pip install llama-index-graph-stores-neptune\n",
    "%pip install llama-index-embeddings-bedrock"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be3f7baa-1c0a-430b-981b-83ddca9e71f2",
   "metadata": {},
   "source": [
    "## Using Knowledge Graph with NeptuneDatabaseGraphStore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97221c15",
   "metadata": {},
   "source": [
    "### Add the required imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79c7f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock import Bedrock\n",
    "from llama_index.embeddings.bedrock import BedrockEmbedding\n",
    "from llama_index.core import (\n",
    "    StorageContext,\n",
    "    SimpleDirectoryReader,\n",
    "    KnowledgeGraphIndex,\n",
    "    Settings,\n",
    ")\n",
    "from llama_index.graph_stores.neptune import (\n",
    "    NeptuneAnalyticsGraphStore,\n",
    "    NeptuneDatabaseGraphStore,\n",
    ")\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f553e01f",
   "metadata": {},
   "source": [
    "### Configure the LLM to use, in this case Amazon Bedrock and Claude 2.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "032264ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = Bedrock(model=\"anthropic.claude-v2\")\n",
    "embed_model = BedrockEmbedding(model=\"amazon.titan-embed-text-v1\")\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model\n",
    "Settings.chunk_size = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75f1d565-04e8-41bc-9165-166dc89b6b47",
   "metadata": {},
   "source": [
    "### Building the Knowledge Graph\n",
    "\n",
    "### Read in the sample file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c297fd3-3424-41d8-9d0d-25fe6310ab62",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\n",
    "    \"../../../../examples/paul_graham_essay/data\"\n",
    ").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0edbc99",
   "metadata": {},
   "source": [
    "### Instantiate Neptune KG Indexes\n",
    "\n",
    "When using Amazon Neptune you can choose to use either Neptune Database or Neptune Analytics.\n",
    "\n",
    "Neptune Database is a serverless graph database designed for optimal scalability and availability. It provides a solution for graph database workloads that need to scale to 100,000 queries per second, Multi-AZ high availability, and multi-Region deployments. You can use Neptune Database for social networking, fraud alerting, and Customer 360 applications.\n",
    "\n",
    "Neptune Analytics is an analytics database engine that can quickly analyze large amounts of graph data in memory to get insights and find trends. Neptune Analytics is a solution for quickly analyzing existing graph databases or graph datasets stored in a data lake. It uses popular graph analytic algorithms and low-latency analytic queries.\n",
    "\n",
    "\n",
    "#### Using Neptune Database\n",
    "If you can choose to use [Neptune Database](https://docs.aws.amazon.com/neptune/latest/userguide/feature-overview.html) to store your KG index you can create the graph store as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ca71c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_store = NeptuneDatabaseGraphStore(\n",
    "    host=\"<GRAPH NAME>.<CLUSTER ID>.<REGION>.neptune.amazonaws.com\", port=8182\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67418411",
   "metadata": {},
   "source": [
    "#### Neptune Analytics\n",
    "\n",
    "If you can choose to use [Neptune Analytics](https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html) to store your KG index you can create the graph store as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6b11a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_store = NeptuneAnalyticsGraphStore(\n",
    "    graph_identifier=\"<INSERT GRAPH IDENIFIER>\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "370fd08f-56ff-4c24-b0c4-c93116a6d482",
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n",
    "\n",
    "# NOTE: can take a while!\n",
    "index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    storage_context=storage_context,\n",
    "    max_triplets_per_chunk=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c39a0eeb-ef16-4982-8ba8-b37c2c5f4437",
   "metadata": {},
   "source": [
    "#### Querying the Knowledge Graph\n",
    "\n",
    "First, we can query and send only the triplets to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670300d8-d0a8-4201-bbcd-4a74b199fcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=False, response_mode=\"tree_summarize\"\n",
    ")\n",
    "\n",
    "response = query_engine.query(\"Tell me more about Interleaf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eecf2d57-3efa-4b0d-941a-95438d42893c",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd32b8e",
   "metadata": {},
   "source": [
    "For more detailed answers, we can also send the text from where the retrieved tripets were extracted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd14686d-1c53-4637-9340-3745f2121ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    include_text=True, response_mode=\"tree_summarize\"\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"Tell me more about what the author worked on at Interleaf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c87d14-d2d8-4d80-89f6-1e5972973528",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}