Add ability to retrieve indices from a graph (run-llama#367)

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

Author: jerryjliu

gpt_index/composability/graph.py

@@ -54,6 +54,16 @@ def _get_default_index_registry() -> IndexRegistry:
     return index_registry
 
 
+def _safe_get_index_struct(
+    docstore: DocumentStore, index_struct_id: str
+) -> IndexStruct:
+    """Try get index struct."""
+    index_struct = docstore.get_document(index_struct_id)
+    if not isinstance(index_struct, IndexStruct):
+        raise ValueError("Invalid `index_struct_id` - must be an IndexStruct")
+    return index_struct
+
+
 class ComposableGraph:
     """Composable graph."""
 
@@ -112,6 +122,18 @@ def query(
         )
         return query_runner.query(query_str, self._index_struct)
 
+    def get_index(
+        self, index_struct_id: str, index_cls: Type[BaseGPTIndex], **kwargs: Any
+    ) -> BaseGPTIndex:
+        """Get index."""
+        index_struct = _safe_get_index_struct(self._docstore, index_struct_id)
+        return index_cls(
+            index_struct=index_struct,
+            docstore=self._docstore,
+            index_registry=self._index_registry,
+            **kwargs
+        )
+
     @classmethod
     def load_from_disk(cls, save_path: str, **kwargs: Any) -> "ComposableGraph":
         """Load index from disk.
@@ -135,9 +157,9 @@ def load_from_disk(cls, save_path: str, **kwargs: Any) -> "ComposableGraph":
             docstore = DocumentStore.load_from_dict(
                 result_dict["docstore"], index_registry.type_to_struct
             )
-            index_struct = docstore.get_document(result_dict["index_struct_id"])
-            if not isinstance(index_struct, IndexStruct):
-                raise ValueError("Invalid `index_struct_id` - must be an IndexStruct")
+            index_struct = _safe_get_index_struct(
+                docstore, result_dict["index_struct_id"]
+            )
             return cls(docstore, index_registry, index_struct, **kwargs)
 
     def save_to_disk(self, save_path: str, **save_kwargs: Any) -> None:

gpt_index/indices/base.py

@@ -93,6 +93,10 @@ def __init__(
         self._index_registry = index_registry or IndexRegistry()
 
         if index_struct is not None:
+            if not isinstance(index_struct, self.index_struct_cls):
+                raise ValueError(
+                    f"index_struct must be of type {self.index_struct_cls}"
+                )
             self._index_struct = index_struct
         else:
             documents = cast(Sequence[DOCUMENTS_INPUT], documents)
@@ -230,7 +234,11 @@ def set_doc_id(self, doc_id: str) -> None:
         If you wish to delete the index struct, you can use this doc_id.
 
         """
+        old_doc_id = self._index_struct.get_doc_id()
         self._index_struct.doc_id = doc_id
+        # Note: we also need to delete old doc_id, and update docstore
+        self._docstore.delete_document(old_doc_id)
+        self._docstore.add_documents([self._index_struct])
 
     def get_doc_id(self) -> str:
         """Get doc_id for index struct.

tests/indices/query/test_recursive.py

@@ -210,6 +210,8 @@ def test_recursive_query_table_list(
     table2 = GPTSimpleKeywordTableIndex(documents[2:3], **table_kwargs)
     table1.set_text("table_summary1")
     table2.set_text("table_summary2")
+    table1.set_doc_id("table1")
+    table2.set_doc_id("table2")
 
     list_index = GPTListIndex([table1, table2], **list_kwargs)
     query_str = "World?"
@@ -232,6 +234,11 @@ def test_recursive_query_table_list(
         response = graph.query(query_str, query_configs=query_configs)
         assert str(response) == ("Test?:This is a test.")
 
+        # test graph.get_index
+        test_table1 = graph.get_index("table1", GPTSimpleKeywordTableIndex)
+        response = test_table1.query("Hello")
+        assert str(response) == ("Hello:Hello world.")
+
 
 @patch.object(TokenTextSplitter, "split_text", side_effect=mock_token_splitter_newline)
 @patch.object(LLMPredictor, "predict", side_effect=mock_llmpredictor_predict)