refactors: index to query mapping, composability (#328)

Co-authored-by: Jerry Liu <jerry@robustintelligence.com>

Author: jerryjliu

docs/how_to/composability.md

@@ -60,6 +60,28 @@ list_index = GPTListIndex([index1, index2, index3])
 ![](/_static/composability/diagram.png)
 
 
+### Defining a Graph Structure
+
+
+Finally, we define a `ComposableGraph` to "wrap" the composed index graph.
+We can do this by simply feeding in the top-level index.
+This wrapper allows us to query, save, and load the graph to/from disk.
+
+```python
+
+from gpt_index.composability import ComposableGraph
+
+graph = ComposableGraph(list_index)
+
+# [Optional] save to disk
+graph.save_to_disk("save_path.json")
+
+# [Optional] load from disk
+graph = ComposableGraph.load_from_disk("save_path.json")
+
+```
+
+
 ### Querying the Top-Level Index
 
 During a query, we would start with the top-level list index. Each node in the list corresponds to an underlying tree index. 
@@ -86,7 +108,7 @@ query_configs = [
     },
     ...
 ]
-response = list_index.query("Where did the author grow up?", mode="recursive", query_configs=query_configs)
+response = graph.query("Where did the author grow up?", query_configs=query_configs)
 ```
 
 ![](/_static/composability/diagram_q1.png)

docs/index.rst

@@ -78,6 +78,7 @@ At the core of GPT Index is a **data structure**. Instead of relying on world kn
 
    reference/indices.rst
    reference/query.rst
+   reference/composability.rst
    reference/readers.rst
    reference/prompts.rst
    reference/llm_predictor.rst

docs/reference/composability.rst

@@ -0,0 +1,10 @@
+.. _Ref-Composability:
+
+Composability
+=============
+
+Below we show the API reference for composable data structures.
+
+.. automodule:: gpt_index.composability
+   :members:
+   :inherited-members:

examples/composable_indices/ComposableIndices.ipynb

@@ -12,12 +12,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "e27b0473-4bda-47f0-b6ed-fd482eac1a13",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gpt_index import GPTTreeIndex, GPTSimpleKeywordTableIndex, GPTListIndex, SimpleDirectoryReader"
+    "from gpt_index import (\n",
+    "    GPTSimpleVectorIndex, \n",
+    "    GPTSimpleKeywordTableIndex, \n",
+    "    GPTListIndex, \n",
+    "    SimpleDirectoryReader\n",
+    ")"
    ]
   },
   {
@@ -102,7 +107,7 @@
    "outputs": [],
    "source": [
     "# build NYC index\n",
-    "nyc_index = GPTTreeIndex(nyc_documents)"
+    "nyc_index = GPTSimpleVectorIndex(nyc_documents)"
    ]
   },
   {
@@ -112,7 +117,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "nyc_index.save_to_disk('../test_wiki/index.json')"
+    "nyc_index.save_to_disk('index_nyc.json')"
    ]
   },
   {
@@ -123,7 +128,7 @@
    "outputs": [],
    "source": [
     "# build essay index\n",
-    "essay_index = GPTTreeIndex(essay_documents)"
+    "essay_index = GPTSimpleVectorIndex(essay_documents)"
    ]
   },
   {
@@ -133,7 +138,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "essay_index.save_to_disk('../paul_graham_essay/index.json')"
+    "essay_index.save_to_disk('index_pg.json')"
    ]
   },
   {
@@ -147,14 +152,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 8,
    "id": "98068ef8-aead-46e7-8dac-0d05b5a86e6a",
    "metadata": {},
    "outputs": [],
    "source": [
     "# try loading\n",
-    "nyc_index = GPTTreeIndex.load_from_disk('../test_wiki/index.json')\n",
-    "essay_index = GPTTreeIndex.load_from_disk('../paul_graham_essay/index.json')"
+    "nyc_index = GPTSimpleVectorIndex.load_from_disk('index_nyc.json')\n",
+    "essay_index = GPTSimpleVectorIndex.load_from_disk('index_pg.json')"
    ]
   },
   {
@@ -171,7 +176,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 9,
    "id": "4149cbbd-7d0b-48c4-8c47-7d67ae0c55f0",
    "metadata": {},
    "outputs": [],
@@ -228,18 +233,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 10,
    "id": "76c251ca-b06b-42e9-ac99-aa0a0a5187d4",
    "metadata": {},
    "outputs": [],
    "source": [
     "# set query config\n",
     "query_configs = [\n",
     "    {\n",
-    "        \"index_struct_type\": \"tree\",\n",
+    "        \"index_struct_type\": \"simple_dict\",\n",
     "        \"query_mode\": \"default\",\n",
     "        \"query_kwargs\": {\n",
-    "            \"child_branch_factor\": 2\n",
+    "            \"similarity_top_k\": 1\n",
     "        }\n",
     "    },\n",
     "    {\n",
@@ -260,6 +265,56 @@
     "keyword_table = GPTSimpleKeywordTableIndex([nyc_index, essay_index], max_keywords_per_chunk=50)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "eebbc448-1e0b-402c-b37e-f93bfcc0bf4f",
+   "metadata": {},
+   "source": [
+    "### Define Graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "6d68750c-e5ae-481a-8b03-6173020c9bf3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from gpt_index.composability import ComposableGraph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "822ada9f-fb43-472e-95ce-0036d508e528",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "graph = ComposableGraph.build_from_index(keyword_table)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "ae127943-afac-48b4-b22d-84a37e553e4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# [optional] save to disk\n",
+    "graph.save_to_disk(\"index_graph.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "dca2b64b-9af1-456f-8dab-822bfdc5d0ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# [optional] load from disk\n",
+    "graph = ComposableGraph.load_from_disk(\"index_graph.json\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -268,24 +323,26 @@
    "outputs": [],
    "source": [
     "# ask it a question about NYC \n",
-    "response = keyword_table.query(\n",
+    "response = graph.query(\n",
     "    \"What is the climate of New York City like? How cold is it during the winter?\", \n",
-    "    mode=\"recursive\", \n",
-    "    query_configs=query_configs\n",
+    "    query_configs=query_configs,\n",
+    "    verbose=True\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 27,
    "id": "c0a43443-3e00-4e48-b3ab-f6369191d53a",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "The climate of New York City is generally mild with hot and humid summers and cool to cold winters. The average temperature in the coldest winter month is 16°F (or -9°C). Nighttime temperatures can be especially cold due to the urban heat island effect. Temperatures can reach 0°F (-18°C) on rare occasions, with the coldest recorded wind chill being -37°F (-38°C). The city receives 49.5 inches (1,260 mm) of precipitation annually, which is relatively evenly spread throughout the year. Average winter snowfall between 1991 and 2020 has been 29.8 inches (76 cm); this varies considerably between years. Hurricanes and tropical storms are rare in the New York area, with the coldest month on record being January 1857, with a mean temperature of 19.6 °F (−6.9 °C). The warmest months on record are July 1825 and July 1999, both with a mean temperature of 81.4 °F (27.4 °C). The warmest years on record are 2012 and 2020, both with mean temperatures of 57.1 °F (13.9 °C). The coldest year is 1836, with a mean temperature of 47.\n"
+      "\n",
+      "\n",
+      "New York City has a humid subtropical climate (Cfa) under the Köppen climate classification. Winters are typically chilly and damp, with temperatures usually dropping to 10 °F (−12 °C) several times per winter, yet can also reach 60 °F (16 °C) for several days even in the coldest winter month. The daily mean temperature in January, the area's coldest month, is 33.3 °F (0.7 °C). The city receives an average of 46.9 inches (1,194 mm) of rainfall annually, with the wettest month being August 2011, with 18.95 inches (481 mm) of rainfall. The snowiest month on record is February 2010, with 36.9 inches (94 cm) of snowfall. The snowiest season (Jul–Jun) on record is 1995–1996, with 75.6 inches (192 cm) of snowfall.\n"
      ]
     }
    ],
@@ -295,25 +352,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 28,
    "id": "c78bc3da-6bad-4998-9a81-90a3fa9200a9",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      ">Source (Doc id: 6eb00cc4-27e9-4ba7-a7b9-bbb2467116a7): \n",
+      "> Source (Doc id: 4e8c9bbc-b42f-479f-8fb1-83d0b6198f1d): \n",
       "    New York, often called New York City or NYC, \n",
       "    is the most populous city in the United St...\n",
       "\n",
-      ">Source (Doc id: d4cf39c8-e10b-4552-9f6d-60bb37356769): °F (16 °C) for several days even in the coldest winter month. Spring and autumn are unpredictable...\n",
-      "\n",
-      ">Source (Doc id: d4cf39c8-e10b-4552-9f6d-60bb37356769): in August.The city receives 49.5 inches (1,260 mm) of precipitation annually, which is relatively...\n",
-      "\n",
-      ">Source (Doc id: d4cf39c8-e10b-4552-9f6d-60bb37356769): the National Hockey League, and Major League Soccer. The New York metropolitan area hosts the mos...\n",
-      "\n",
-      ">Source (Doc id: d4cf39c8-e10b-4552-9f6d-60bb37356769): any city in North America. New York City is the host of Climate Week NYC, the largest Climate Wee...\n"
+      "> Source (Doc id: 77f3b3ea-93ab-49c8-b938-2bd3c870a602): has been altered substantially by human intervention, with considerable land reclamation along th...\n"
      ]
     }
    ],
@@ -330,28 +381,26 @@
    "outputs": [],
    "source": [
     "# ask it a question about PG's essay\n",
-    "response = keyword_table.query(\n",
+    "response = graph.query(\n",
     "    \"What did the author do growing up, before his time at Y Combinator?\", \n",
-    "    mode=\"recursive\", \n",
     "    query_configs=query_configs\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 30,
    "id": "06dc71bb-882d-49f5-8566-69b0ea5019dd",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "'The author was a writer and programmer who studied art and worked on building a WYSIWYG site builder.'"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "The author grew up in England and attended college in the United States. He studied computer science and art, and worked on a variety of projects, including writing essays, hacking, and working on a Lisp interpreter. He also worked on a startup called Viaweb, which was eventually acquired by Yahoo. He also worked on Interleaf, a high-end, special-purpose hardware and software company, and sought out signature styles at RISD. He also lived in a rent-stabilized apartment in New York, and worked on software projects that could be launched as soon as they were done.\n"
+     ]
     }
    ],
    "source": [
@@ -360,10 +409,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 31,
    "id": "b0894565-2b2c-4987-a891-17ba44d775b5",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "> Source (Doc id: ae92ab9a-c6ed-48c3-b333-d459908dec3f): \n",
+      "    Author: Paul Graham. \n",
+      "    The author grew up painting and writing essays. \n",
+      "    He wrote a bo...\n",
+      "\n",
+      "> Source (Doc id: a63ad7c4-87f2-42fd-a32b-f682a022af90): get their initial set of customers almost entirely from among their batchmates.\n",
+      "\n",
+      "I had not origin...\n"
+     ]
+    }
+   ],
    "source": [
     "# Get source of response\n",
     "print(response.get_formatted_sources())"

examples/composable_indices/index_graph.json

@@ -0,0 +1,6 @@
+"""This module contains all classes used for composing graphs over indices."""
+
+
+from gpt_index.composability.graph import ComposableGraph
+
+__all__ = ["ComposableGraph"]