feat: add graph_stores, impl Simple KG & Nebula KG

run-llama · May 25, 2023 · bfb3e79 · bfb3e79
1 parent 429260f
commit bfb3e79
Show file tree

Hide file tree

Showing 19 changed files with 1,483 additions and 161 deletions.
diff --git a/docs/examples/composable_indices/city_analysis/PineconeDemo-CityAnalysis.ipynb b/docs/examples/composable_indices/city_analysis/PineconeDemo-CityAnalysis.ipynb
@@ -1,6 +1,7 @@
 {
     "cells": [
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "cfb64210-9c6b-47d7-81f4-67dbdab68e4c",
             "metadata": {
@@ -26,7 +27,7 @@
             },
             "outputs": [],
             "source": [
-                "!pip install llama-index"
+                "%pip install llama-index"
             ]
         },
         {
@@ -82,6 +83,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "49e0d841-680f-4a0c-b455-788b54978ebf",
             "metadata": {
@@ -160,6 +162,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "84bfcaa1-db15-45ba-8af1-fee548354965",
             "metadata": {},
@@ -221,6 +224,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "b373db78",
             "metadata": {},
@@ -258,6 +262,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "b1d69b03",
             "metadata": {},
@@ -315,6 +320,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "17605939-09ce-4405-a92a-8f296c941893",
             "metadata": {},
@@ -350,6 +356,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "d4d3cd8b-4134-4cfa-8002-e0a34694d2e1",
             "metadata": {
@@ -426,6 +433,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "49c900ee-a31f-4fcd-bb44-ff2cd12a41eb",
             "metadata": {
@@ -436,6 +444,7 @@
             ]
         },
         {
+            "attachments": {},
             "cell_type": "markdown",
             "id": "e0a8fa6a-e96e-4341-bb43-7547415f766e",
             "metadata": {
@@ -1883,4 +1892,4 @@
     },
     "nbformat": 4,
     "nbformat_minor": 5
-}
+}
diff --git a/docs/examples/index_structs/knowledge_graph/KnowledgeGraphDemo.ipynb b/docs/examples/index_structs/knowledge_graph/KnowledgeGraphDemo.ipynb
@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "82f90261",
    "metadata": {},
@@ -34,6 +33,16 @@
     "logging.basicConfig(stream=sys.stdout, level=logging.INFO)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "be3f7baa-1c0a-430b-981b-83ddca9e71f2",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Using Knowledge Graph"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "75f1d565-04e8-41bc-9165-166dc89b6b47",
@@ -49,8 +58,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from llama_index import SimpleDirectoryReader, LLMPredictor, ServiceContext\n",
-    "from llama_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex\n",
+    "from llama_index import SimpleDirectoryReader, LLMPredictor, ServiceContext, GPTKnowledgeGraphIndex\n",
+    "from llama_index.graph_stores import NebulaGraphStore\n",
+    "\n",
     "from langchain import OpenAI\n",
     "from IPython.display import Markdown, display"
    ]
@@ -76,6 +86,7 @@
    "source": [
     "# define LLM\n",
     "# NOTE: at the time of demo, text-davinci-002 did not have rate-limit errors\n",
+    "\n",
     "llm_predictor = LLMPredictor(llm=OpenAI(temperature=0, model_name=\"text-davinci-002\"))\n",
     "service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, chunk_size_limit=512)"
    ]
@@ -89,12 +100,18 @@
    },
    "outputs": [],
    "source": [
+    "from llama_index.storage.storage_context import StorageContext\n",
+    "\n",
+    "graph_store = NebulaGraphStore()\n",
+    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n",
+    "\n",
     "# NOTE: can take a while! \n",
-    "new_index = GPTKnowledgeGraphIndex.from_documents(\n",
-    "    documents, \n",
+    "index = GPTKnowledgeGraphIndex.from_documents(\n",
+    "    documents,\n",
     "    max_triplets_per_chunk=2,\n",
+    "    storage_context=storage_context,\n",
     "    service_context=service_context\n",
-    ")"
+    ")\n"
    ]
   },
   {
@@ -137,7 +154,7 @@
     }
    ],
    "source": [
-    "query_engine = new_index.as_query_engine(\n",
+    "query_engine = index.as_query_engine(\n",
     "    include_text=False, \n",
     "    response_mode=\"tree_summarize\"\n",
     ")\n",
@@ -195,7 +212,7 @@
     }
    ],
    "source": [
-    "query_engine = new_index.as_query_engine(\n",
+    "query_engine = index.as_query_engine(\n",
     "    include_text=True, \n",
     "    response_mode=\"tree_summarize\"\n",
     ")\n",
@@ -253,7 +270,7 @@
    ],
    "source": [
     "# NOTE: can take a while! \n",
-    "new_index = GPTKnowledgeGraphIndex.from_documents(\n",
+    "index = GPTSimpleKGIndex.from_documents(\n",
     "    documents, \n",
     "    max_triplets_per_chunk=2,\n",
     "    service_context=service_context,\n",
@@ -309,7 +326,7 @@
    ],
    "source": [
     "# query using top 3 triplets plus keywords (duplicate triplets are removed)\n",
-    "query_engine = new_index.as_query_engine(\n",
+    "query_engine = index.as_query_engine(\n",
     "    include_text=True, \n",
     "    response_mode=\"tree_summarize\",\n",
     "    embedding_mode='hybrid',\n",
@@ -385,7 +402,7 @@
     "## create graph\n",
     "from pyvis.network import Network\n",
     "\n",
-    "g = new_index.get_networkx_graph()\n",
+    "g = index.get_networkx_graph()\n",
     "net = Network(notebook=True, cdn_resources=\"in_line\", directed=True)\n",
     "net.from_nx(g)\n",
     "net.show(\"example.html\")"
@@ -497,9 +514,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "llama_index",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "llama_index"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -511,7 +528,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.10"
+   "version": "3.9.16"
   }
  },
  "nbformat": 4,