Add start and end indices to Node extra info during index construction

gpt_index/data_structs/data_structs.py

@@ -3,7 +3,7 @@
 import random
 import sys
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional, Set
 
 from dataclasses_json import DataClassJsonMixin
 
@@ -51,6 +51,9 @@ def __post_init__(self) -> None:
     # reference document id
     ref_doc_id: Optional[str] = None
 
+    # extra node info
+    node_info: Optional[Dict[str, Any]] = None
+
     def get_text(self) -> str:
         """Get text."""
         text = super().get_text()

gpt_index/indices/base.py

@@ -196,16 +196,23 @@ def _get_nodes_from_document(
         """Add document to index."""
         text_chunks = text_splitter.split_text(document.get_text())
         nodes = []
+        index_counter = 0
         for i, text_chunk in enumerate(text_chunks):
             fmt_text_chunk = truncate_text(text_chunk, 50)
             print(f"> Adding chunk: {fmt_text_chunk}")
+            index_pos_info = {
+                "start": index_counter,  # NOTE: start is inclusive
+                "end": index_counter + len(text_chunk),  # NOTE: end is exclusive
+            }
+            index_counter += len(text_chunk) + 1
             # if embedding specified in document, pass it to the Node
             node = Node(
                 text=text_chunk,
                 index=start_idx + i,
                 ref_doc_id=document.get_doc_id(),
                 embedding=document.embedding,
                 extra_info=document.extra_info if self._include_extra_info else None,
+                node_info=index_pos_info,
             )
             nodes.append(node)
         return nodes

tests/indices/list/test_base.py

@@ -257,3 +257,21 @@ def test_extra_info(
     assert list_index.index_struct.nodes[3].get_text() == (
         "extra_info: extra_info\n" "foo: bar\n\n" "This is a test v2."
     )
+
+
+@patch_common
+def test_node_indices(
+    _mock_init: Any,
+    _mock_predict: Any,
+    _mock_total_tokens_used: Any,
+    _mock_split_text: Any,
+    documents: List[Document],
+    struct_kwargs: Dict,
+) -> None:
+    """Test Node start and end indices info."""
+    index_kwargs, query_kwargs = struct_kwargs
+    index = GPTListIndex(documents, **index_kwargs)
+
+    assert index.index_struct.nodes[0].node_info == {"start": 0, "end": 12}
+    assert index.index_struct.nodes[1].node_info == {"start": 13, "end": 28}
+    assert index.index_struct.nodes[2].node_info == {"start": 29, "end": 50}