feat: Knowledge Graph Query Engine (#6642)

run-llama · Jul 12, 2023 · 93f5f28 · 93f5f28
1 parent 582686b
commit 93f5f28
Show file tree

Hide file tree

Showing 18 changed files with 2,763 additions and 136 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 ### New Features
 - Improved milvus consistency support and output fields support (#6452)
+- Added support for knowledge graph querying w/ cypyer+nebula (#6642)
 
 ## [v0.7.6] - 2023-07-12
 

diff --git a/docs/api_reference/query/query_engines.rst b/docs/api_reference/query/query_engines.rst
@@ -19,6 +19,7 @@ Below we show some general query engine classes.
    query_engines/sql_join_query_engine.rst
    query_engines/flare_query_engine.rst
    query_engines/citation_query_engine.rst
+   query_engines/knowledge_graph_query_engine.rst
 
 
 We also show query engine classes specific to our structured indices.

diff --git a/docs/api_reference/query/query_engines/knowledge_graph_query_engine.rst b/docs/api_reference/query/query_engines/knowledge_graph_query_engine.rst
@@ -0,0 +1,6 @@
+Knowledge Graph Query Engine
+============================
+
+.. automodule:: llama_index.query_engine.knowledge_graph_query_engine
+   :members:
+   :inherited-members:
diff --git a/docs/community/integrations.md b/docs/community/integrations.md
@@ -11,5 +11,6 @@ integrations/guidance.md
 integrations/trulens.md
 integrations/chatgpt_plugins.md
 integrations/using_with_langchain.md
+integrations/graph_stores.md
 integrations/vector_stores.md
 ```
diff --git a/docs/community/integrations/graph_stores.md b/docs/community/integrations/graph_stores.md
@@ -0,0 +1,15 @@
+# Using Graph Stores
+
+## `NebulaGraphStore`
+
+We support a `NebulaGraphStore` integration, for persisting graphs directly in Nebula! Furthermore, you can generate cypher queries and return natural language responses for your Nebula graphs using the `KnowledgeGraphQueryEngine`.
+
+See the associated guides below:
+
+```{toctree}
+---
+maxdepth: 1
+---
+Nebula Graph Store </examples/index_structs/knowledge_graph/NebulaGraphKGIndexDemo.ipynb>
+Knowledge Graph Query Engine </examples/query_engine/knowledge_graph_query_engine.ipynb>
+```
diff --git a/docs/core_modules/query_modules/query_engine/modules.md b/docs/core_modules/query_modules/query_engine/modules.md
@@ -16,6 +16,7 @@ maxdepth: 1
 ---
 /examples/query_engine/json_query_engine.ipynb
 /examples/query_engine/pandas_query_engine.ipynb
+/examples/query_engine/knowledge_graph_query_engine.ipynb
 ```
 
 ## Advanced

diff --git a/...cts/knowledge_graph/KnowledgeGraphIndex_vs_VectorStoreIndex_vs_CustomIndex_combined.ipynb b/...cts/knowledge_graph/KnowledgeGraphIndex_vs_VectorStoreIndex_vs_CustomIndex_combined.ipynb
@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "105a2123",
    "metadata": {},
@@ -132,8 +131,6 @@
    "outputs": [],
    "source": [
     "# For Azure OpenAI\n",
-    "%load_ext autoreload\n",
-    "%autoreload 2\n",
     "\n",
     "import os\n",
     "import json\n",
@@ -171,13 +168,20 @@
     "llm = AzureOpenAI(\n",
     "    engine=\"<foo-bar-deployment>\",\n",
     "    temperature=0,\n",
+    "    openai_api_version=openai.api_version,\n",
+    "    model_kwargs={\n",
+    "        \"api_key\": openai.api_key,\n",
+    "        \"api_base\": openai.api_base,\n",
+    "        \"api_type\": openai.api_type,\n",
+    "        \"api_version\": openai.api_version,\n",
+    "    },\n",
     ")\n",
     "\n",
     "# You need to deploy your own embedding model as well as your own chat completion model\n",
     "embedding_llm = LangchainEmbedding(\n",
     "    OpenAIEmbeddings(\n",
     "        model=\"text-embedding-ada-002\",\n",
-    "        deployment=\"nebula_docs_2_embedding\",\n",
+    "        deployment=\"<foo-bar-deployment>\",\n",
     "        openai_api_key=openai.api_key,\n",
     "        openai_api_base=openai.api_base,\n",
     "        openai_api_type=openai.api_type,\n",
@@ -195,7 +199,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "2b17442e",
    "metadata": {},
@@ -216,29 +219,29 @@
     "os.environ[\"NEBULA_PASSWORD\"] = \"nebula\"\n",
     "os.environ[\n",
     "    \"NEBULA_ADDRESS\"\n",
-    "] = \"127.0.0.1:9669\"  # assumed we have NebulaGraph installed locally\n",
+    "] = \"127.0.0.1:9669\"  # assumed we have NebulaGraph 3.5.0 or newer installed locally\n",
     "\n",
     "# Assume that the graph has already been created\n",
     "# Create a NebulaGraph cluster with:\n",
     "# Option 0: `curl -fsSL nebula-up.siwei.io/install.sh | bash`\n",
     "# Option 1: NebulaGraph Docker Extension https://hub.docker.com/extensions/weygu/nebulagraph-dd-ext\n",
-    "# and that the graph space is called \"test\"\n",
+    "# and that the graph space is called \"llamaindex\"\n",
     "# If not, create it with the following commands from NebulaGraph's console:\n",
     "# CREATE SPACE llamaindex(vid_type=FIXED_STRING(256), partition_num=1, replica_factor=1);\n",
     "# :sleep 10;\n",
     "# USE llamaindex;\n",
-    "# CREATE TAG entity();\n",
-    "# CREATE EDGE rel(predicate string);\n",
+    "# CREATE TAG entity(name string);\n",
+    "# CREATE EDGE relationship(relationship string);\n",
+    "# CREATE TAG INDEX entity_index ON entity(name(256));\n",
     "\n",
     "space_name = \"llamaindex\"\n",
-    "edge_types, rel_prop_names = [\"rel\"], [\n",
-    "    \"predicate\"\n",
+    "edge_types, rel_prop_names = [\"relationship\"], [\n",
+    "    \"relationship\"\n",
     "]  # default, could be omit if create from an empty kg\n",
     "tags = [\"entity\"]  # default, could be omit if create from an empty kg"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "d9cb7083",
    "metadata": {},
@@ -263,7 +266,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "cc21c03b",
    "metadata": {},
@@ -312,7 +314,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "3c3efdc4",
    "metadata": {},
@@ -344,7 +345,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "7ac0ee74",
    "metadata": {},
@@ -414,7 +414,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "fbe22b2b",
    "metadata": {},
@@ -447,7 +446,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "541608be",
    "metadata": {},
@@ -480,7 +478,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "d6cd7e11",
    "metadata": {},
@@ -557,7 +554,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "adb2311e",
    "metadata": {},
@@ -613,7 +609,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "2ac8c99a",
    "metadata": {},
@@ -703,7 +698,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "bb2eb936",
    "metadata": {},
@@ -731,7 +725,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "7ecc23e9",
    "metadata": {},
@@ -922,7 +915,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "41ee9f9b",
    "metadata": {},
@@ -1010,9 +1002,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.10"
+   "version": "3.11.4"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}