diff --git a/docs/guides/primer/usage_pattern.md b/docs/guides/primer/usage_pattern.md index 4d40d08bb150b..bdd30446aebd2 100644 --- a/docs/guides/primer/usage_pattern.md +++ b/docs/guides/primer/usage_pattern.md @@ -54,7 +54,7 @@ nodes = parser.get_nodes_from_documents(documents) You can also choose to construct Node objects manually and skip the first section. For instance, ```python -from llama_index.data_structs.node_v2 import Node, DocumentRelationship +from llama_index.data_structs.node import Node, DocumentRelationship node1 = Node(text="", doc_id="") node2 = Node(text="", doc_id="") @@ -263,6 +263,7 @@ from llama_index import ( ) from llama_index.retrievers import VectorIndexRetriever from llama_index.query_engine import RetrieverQueryEngine +from llama_index.indices.postprocessor import SimilarityPostprocessor # build index index = GPTVectorStoreIndex.from_documents(documents) @@ -276,7 +277,7 @@ retriever = VectorIndexRetriever( # configure response synthesizer response_synthesizer = ResponseSynthesizer.from_args( node_postprocessors=[ - SimilarityPostprocessor(required_keywords=['hello']) + SimilarityPostprocessor(similarity_cutoff=0.7) ] ) @@ -285,6 +286,10 @@ query_engine = RetrieverQueryEngine( retriever=retriever, response_synthesizer=response_synthesizer, ) + +# query +response = query_engine.query("What did the author do growing up?") +print(response) ``` You may also add your own retrieval, response synthesis, and overall query logic, by implementing the corresponding interfaces. diff --git a/docs/reference/node.rst b/docs/reference/node.rst index 032be6bc86498..da1e77b4be5fa 100644 --- a/docs/reference/node.rst +++ b/docs/reference/node.rst @@ -3,7 +3,7 @@ Node ================= -.. automodule:: llama_index.data_structs.node_v2 +.. automodule:: llama_index.data_structs.node :members: :inherited-members: :exclude-members: NodeType, ImageNode, IndexNode \ No newline at end of file diff --git a/examples/composable_indices/City_Analysis-Decompose-KeywordTable.ipynb b/examples/composable_indices/City_Analysis-Decompose-KeywordTable.ipynb index aa0a26ce88587..01901ac3841f7 100644 --- a/examples/composable_indices/City_Analysis-Decompose-KeywordTable.ipynb +++ b/examples/composable_indices/City_Analysis-Decompose-KeywordTable.ipynb @@ -763,7 +763,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "/Users/suo/dev/llama_index/llama_index/data_structs/node_v2.py:176: UserWarning: .source_text is deprecated, use .node.get_text() instead\n", + "/Users/suo/dev/llama_index/llama_index/data_structs/node.py:176: UserWarning: .source_text is deprecated, use .node.get_text() instead\n", " warnings.warn(\".source_text is deprecated, use .node.get_text() instead\")\n" ] } diff --git a/llama_index/data_structs/__init__.py b/llama_index/data_structs/__init__.py index 507295a07e03a..b6c6d52f7a00f 100644 --- a/llama_index/data_structs/__init__.py +++ b/llama_index/data_structs/__init__.py @@ -1,13 +1,13 @@ """Init file.""" -from llama_index.data_structs.data_structs_v2 import ( +from llama_index.data_structs.data_structs import ( IndexDict, IndexGraph, IndexList, KeywordTable, ) -from llama_index.data_structs.node_v2 import Node, NodeWithScore -from llama_index.data_structs.table_v2 import StructDatapoint +from llama_index.data_structs.node import Node, NodeWithScore +from llama_index.data_structs.table import StructDatapoint __all__ = [ "Node", diff --git a/llama_index/data_structs/data_structs.py b/llama_index/data_structs/data_structs.py index 642ba87cfae89..e518041fdce56 100644 --- a/llama_index/data_structs/data_structs.py +++ b/llama_index/data_structs/data_structs.py @@ -1,149 +1,130 @@ -"""File for core data structures.""" +"""Data structures. -import random -import sys +Nodes are decoupled from the indices. + +""" + +import uuid +from abc import abstractmethod from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Dict, List, Optional, Sequence, Set, Tuple from dataclasses_json import DataClassJsonMixin +from llama_index.data_structs.node import Node from llama_index.data_structs.struct_type import IndexStructType -from llama_index.schema import BaseDocument -from llama_index.utils import get_new_int_id @dataclass -class IndexStruct(BaseDocument, DataClassJsonMixin): +class IndexStruct(DataClassJsonMixin): """A base data struct for a LlamaIndex.""" - # NOTE: the text field, inherited from BaseDocument, - # represents a summary of the content of the index struct. - # primarily used for composing indices with other indices - - # NOTE: the doc_id field, inherited from BaseDocument, - # represents a unique identifier for the index struct - # that will be put in the docstore. - # Not all index_structs need to have a doc_id. Only index_structs that - # represent a complete data structure (e.g. IndexGraph, IndexList), - # and are used to compose a higher level index, will have a doc_id. - - -@dataclass -class Node(IndexStruct): - """A generic node of data. - - Base struct used in most indices. - - """ - - def __post_init__(self) -> None: - """Post init.""" - super().__post_init__() - # NOTE: for Node objects, the text field is required - if self.text is None: - raise ValueError("text field not set.") - - # used for GPTTreeIndex - index: int = 0 - child_indices: Set[int] = field(default_factory=set) + index_id: str = field(default_factory=lambda: str(uuid.uuid4())) + summary: Optional[str] = None - # embeddings - embedding: Optional[List[float]] = None - - # reference document id - ref_doc_id: Optional[str] = None - - # extra node info - node_info: Optional[Dict[str, Any]] = None - - # TODO: store reference instead of actual image - # base64 encoded image str - image: Optional[str] = None - - def get_text(self) -> str: - """Get text.""" - text = super().get_text() - result_text = ( - text if self.extra_info_str is None else f"{self.extra_info_str}\n\n{text}" - ) - return result_text + def get_summary(self) -> str: + """Get text summary.""" + if self.summary is None: + raise ValueError("summary field of the index_struct not set.") + return self.summary @classmethod - def get_type(cls) -> str: - """Get type.""" - # TODO: consolidate with IndexStructType - return "node" + @abstractmethod + def get_type(cls) -> IndexStructType: + """Get index struct type.""" @dataclass class IndexGraph(IndexStruct): """A graph representing the tree-structured index.""" - all_nodes: Dict[int, Node] = field(default_factory=dict) - root_nodes: Dict[int, Node] = field(default_factory=dict) + # mapping from index in tree to Node doc id. + all_nodes: Dict[int, str] = field(default_factory=dict) + root_nodes: Dict[int, str] = field(default_factory=dict) + node_id_to_children_ids: Dict[str, List[str]] = field(default_factory=dict) + + @property + def node_id_to_index(self) -> Dict[str, int]: + """Map from node id to index.""" + return {node_id: index for index, node_id in self.all_nodes.items()} @property def size(self) -> int: """Get the size of the graph.""" return len(self.all_nodes) - def get_children(self, parent_node: Optional[Node]) -> Dict[int, Node]: - """Get nodes given indices.""" + def get_index(self, node: Node) -> int: + """Get index of node.""" + return self.node_id_to_index[node.get_doc_id()] + + def insert( + self, + node: Node, + index: Optional[int] = None, + children_nodes: Optional[Sequence[Node]] = None, + ) -> None: + """Insert node.""" + index = index or self.size + node_id = node.get_doc_id() + + self.all_nodes[index] = node_id + + if children_nodes is None: + children_nodes = [] + children_ids = [n.get_doc_id() for n in children_nodes] + self.node_id_to_children_ids[node_id] = children_ids + + def get_children(self, parent_node: Optional[Node]) -> Dict[int, str]: + """Get children nodes.""" if parent_node is None: return self.root_nodes else: - return {i: self.all_nodes[i] for i in parent_node.child_indices} - - def insert_under_parent(self, node: Node, parent_node: Optional[Node]) -> None: + parent_id = parent_node.get_doc_id() + children_ids = self.node_id_to_children_ids[parent_id] + return { + self.node_id_to_index[child_id]: child_id for child_id in children_ids + } + + def insert_under_parent( + self, node: Node, parent_node: Optional[Node], new_index: Optional[int] = None + ) -> None: """Insert under parent node.""" - if node.index in self.all_nodes: - raise ValueError( - "Cannot insert a new node with the same index as an existing node." - ) + new_index = new_index or self.size if parent_node is None: - self.root_nodes[node.index] = node + self.root_nodes[new_index] = node.get_doc_id() + self.node_id_to_children_ids[node.get_doc_id()] = [] else: - parent_node.child_indices.add(node.index) + if parent_node.doc_id not in self.node_id_to_children_ids: + self.node_id_to_children_ids[parent_node.get_doc_id()] = [] + self.node_id_to_children_ids[parent_node.get_doc_id()].append( + node.get_doc_id() + ) - self.all_nodes[node.index] = node + self.all_nodes[new_index] = node.get_doc_id() @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" - return "tree" + return IndexStructType.TREE @dataclass class KeywordTable(IndexStruct): """A table of keywords mapping keywords to text chunks.""" - table: Dict[str, Set[int]] = field(default_factory=dict) - text_chunks: Dict[int, Node] = field(default_factory=dict) - - def _get_index(self) -> int: - """Get the next index for the text chunk.""" - # randomly generate until we get a unique index - while True: - idx = random.randint(0, sys.maxsize) - if idx not in self.text_chunks: - break - return idx + table: Dict[str, Set[str]] = field(default_factory=dict) - def add_node(self, keywords: List[str], node: Node) -> int: + def add_node(self, keywords: List[str], node: Node) -> None: """Add text to table.""" - cur_idx = self._get_index() for keyword in keywords: if keyword not in self.table: self.table[keyword] = set() - self.table[keyword].add(cur_idx) - self.text_chunks[cur_idx] = node - return cur_idx + self.table[keyword].add(node.get_doc_id()) - def get_texts(self, keyword: str) -> List[str]: - """Get texts given keyword.""" - if keyword not in self.table: - raise ValueError("Keyword not found in table.") - return [self.text_chunks[idx].get_text() for idx in self.table[keyword]] + @property + def node_ids(self) -> Set[str]: + """Get all node ids.""" + return set.union(*self.table.values()) @property def keywords(self) -> Set[str]: @@ -156,34 +137,36 @@ def size(self) -> int: return len(self.table) @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" - return "keyword_table" + return IndexStructType.KEYWORD_TABLE @dataclass class IndexList(IndexStruct): """A list of documents.""" - nodes: List[Node] = field(default_factory=list) + nodes: List[str] = field(default_factory=list) def add_node(self, node: Node) -> None: """Add text to table, return current position in list.""" # don't worry about child indices for now, nodes are all in order - self.nodes.append(node) + self.nodes.append(node.get_doc_id()) @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" - return "list" + return IndexStructType.LIST @dataclass class IndexDict(IndexStruct): """A simple dictionary of documents.""" - nodes_dict: Dict[int, Node] = field(default_factory=dict) - id_map: Dict[str, int] = field(default_factory=dict) + # mapping from vector store id to node id + nodes_dict: Dict[str, str] = field(default_factory=dict) + # mapping from doc_id to vector store id + doc_id_dict: Dict[str, List[str]] = field(default_factory=dict) # TODO: temporary hack to store embeddings for simple vector index # this should be empty for all other indices @@ -195,54 +178,27 @@ def add_node( text_id: Optional[str] = None, ) -> str: """Add text to table, return current position in list.""" - int_id = get_new_int_id(set(self.nodes_dict.keys())) - if text_id in self.id_map: - raise ValueError("text_id cannot already exist in index.") - elif text_id is not None and not isinstance(text_id, str): - raise ValueError("text_id must be a string.") - elif text_id is None: - text_id = str(int_id) - self.id_map[text_id] = int_id + # # don't worry about child indices for now, nodes are all in order + # self.nodes_dict[int_id] = node + vector_id = text_id if text_id is not None else node.get_doc_id() + self.nodes_dict[vector_id] = node.get_doc_id() + if node.ref_doc_id is not None: + if node.ref_doc_id not in self.doc_id_dict: + self.doc_id_dict[node.ref_doc_id] = [] + self.doc_id_dict[node.ref_doc_id].append(vector_id) - # don't worry about child indices for now, nodes are all in order - self.nodes_dict[int_id] = node - return text_id - - def get_nodes(self, text_ids: List[str]) -> List[Node]: - """Get nodes.""" - nodes = [] - for text_id in text_ids: - if text_id not in self.id_map: - raise ValueError("text_id not found in id_map") - elif not isinstance(text_id, str): - raise ValueError("text_id must be a string.") - int_id = self.id_map[text_id] - if int_id not in self.nodes_dict: - raise ValueError("int_id not found in nodes_dict") - nodes.append(self.nodes_dict[int_id]) - return nodes - - def get_node(self, text_id: str) -> Node: - """Get node.""" - return self.get_nodes([text_id])[0] + return vector_id def delete(self, doc_id: str) -> None: - """Delete a document.""" - text_ids_to_delete = set() - int_ids_to_delete = set() - for text_id, int_id in self.id_map.items(): - node = self.nodes_dict[int_id] - if node.ref_doc_id != doc_id: - continue - text_ids_to_delete.add(text_id) - int_ids_to_delete.add(int_id) - - for int_id, text_id in zip(int_ids_to_delete, text_ids_to_delete): - del self.nodes_dict[int_id] - del self.id_map[text_id] + """Delete a Node.""" + if doc_id not in self.doc_id_dict: + return + for vector_id in self.doc_id_dict[doc_id]: + del self.nodes_dict[vector_id] + del self.doc_id_dict[doc_id] @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" return IndexStructType.VECTOR_STORE @@ -254,18 +210,22 @@ class KG(IndexStruct): # Unidirectional table: Dict[str, Set[str]] = field(default_factory=dict) - text_chunks: Dict[str, Node] = field(default_factory=dict) + # text_chunks: Dict[str, Node] = field(default_factory=dict) rel_map: Dict[str, List[Tuple[str, str]]] = field(default_factory=dict) embedding_dict: Dict[str, List[float]] = field(default_factory=dict) + @property + def node_ids(self) -> Set[str]: + """Get all node ids.""" + return set.union(*self.table.values()) + def add_to_embedding_dict(self, triplet_str: str, embedding: List[float]) -> None: """Add embedding to dict.""" self.embedding_dict[triplet_str] = embedding - def upsert_triplet(self, triplet: Tuple[str, str, str], node: Node) -> None: + def upsert_triplet(self, triplet: Tuple[str, str, str]) -> None: """Upsert a knowledge triplet to the graph.""" subj, relationship, obj = triplet - self.add_node([subj, obj], node) if subj not in self.rel_map: self.rel_map[subj] = [] self.rel_map[subj].append((obj, relationship)) @@ -277,7 +237,7 @@ def add_node(self, keywords: List[str], node: Node) -> None: if keyword not in self.table: self.table[keyword] = set() self.table[keyword].add(node_id) - self.text_chunks[node_id] = node + # self.text_chunks[node_id] = node def get_rel_map_texts(self, keyword: str) -> List[str]: """Get the corresponding knowledge for a given keyword.""" @@ -315,90 +275,16 @@ def get_node_ids(self, keyword: str, depth: int = 1) -> List[str]: return node_ids @classmethod - def get_type(cls) -> str: - """Get type.""" - return "kg" - - -# TODO: remove once we centralize UX around vector index - - -class SimpleIndexDict(IndexDict): - """Index dict for simple vector index.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.SIMPLE_DICT - - -class FaissIndexDict(IndexDict): - """Index dict for Faiss vector index.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.DICT - - -class WeaviateIndexDict(IndexDict): - """Index dict for Weaviate vector index.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.WEAVIATE - - -class PineconeIndexDict(IndexDict): - """Index dict for Pinecone vector index.""" - - @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" - return IndexStructType.PINECONE + return IndexStructType.KG -class QdrantIndexDict(IndexDict): - """Index dict for Qdrant vector index.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.QDRANT - - -class ChromaIndexDict(IndexDict): - """Index dict for Chroma vector index.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.CHROMA - - -class OpensearchIndexDict(IndexDict): - """Index dict for Opensearch vector index.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.OPENSEARCH - - -class ChatGPTRetrievalPluginIndexDict(IndexDict): - """Index dict for ChatGPT Retrieval Plugin.""" - - @classmethod - def get_type(cls) -> str: - """Get type.""" - return IndexStructType.CHATGPT_RETRIEVAL_PLUGIN - - -class EmptyIndex(IndexStruct): +@dataclass +class EmptyIndex(IndexDict): """Empty index.""" @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" return IndexStructType.EMPTY diff --git a/llama_index/data_structs/data_structs_v2.py b/llama_index/data_structs/data_structs_v2.py deleted file mode 100644 index a3c3ce5b43f6e..0000000000000 --- a/llama_index/data_structs/data_structs_v2.py +++ /dev/null @@ -1,290 +0,0 @@ -"""Data structures v2. - -Nodes are decoupled from the indices. - -""" - -import uuid -from abc import abstractmethod -from dataclasses import dataclass, field -from typing import Dict, List, Optional, Sequence, Set, Tuple - -from dataclasses_json import DataClassJsonMixin - -from llama_index.data_structs.node_v2 import Node -from llama_index.data_structs.struct_type import IndexStructType - - -@dataclass -class V2IndexStruct(DataClassJsonMixin): - """A base data struct for a LlamaIndex.""" - - index_id: str = field(default_factory=lambda: str(uuid.uuid4())) - summary: Optional[str] = None - - def get_summary(self) -> str: - """Get text summary.""" - if self.summary is None: - raise ValueError("summary field of the index_struct not set.") - return self.summary - - @classmethod - @abstractmethod - def get_type(cls) -> IndexStructType: - """Get index struct type.""" - - -@dataclass -class IndexGraph(V2IndexStruct): - """A graph representing the tree-structured index.""" - - # mapping from index in tree to Node doc id. - all_nodes: Dict[int, str] = field(default_factory=dict) - root_nodes: Dict[int, str] = field(default_factory=dict) - node_id_to_children_ids: Dict[str, List[str]] = field(default_factory=dict) - - @property - def node_id_to_index(self) -> Dict[str, int]: - """Map from node id to index.""" - return {node_id: index for index, node_id in self.all_nodes.items()} - - @property - def size(self) -> int: - """Get the size of the graph.""" - return len(self.all_nodes) - - def get_index(self, node: Node) -> int: - """Get index of node.""" - return self.node_id_to_index[node.get_doc_id()] - - def insert( - self, - node: Node, - index: Optional[int] = None, - children_nodes: Optional[Sequence[Node]] = None, - ) -> None: - """Insert node.""" - index = index or self.size - node_id = node.get_doc_id() - - self.all_nodes[index] = node_id - - if children_nodes is None: - children_nodes = [] - children_ids = [n.get_doc_id() for n in children_nodes] - self.node_id_to_children_ids[node_id] = children_ids - - def get_children(self, parent_node: Optional[Node]) -> Dict[int, str]: - """Get children nodes.""" - if parent_node is None: - return self.root_nodes - else: - parent_id = parent_node.get_doc_id() - children_ids = self.node_id_to_children_ids[parent_id] - return { - self.node_id_to_index[child_id]: child_id for child_id in children_ids - } - - def insert_under_parent( - self, node: Node, parent_node: Optional[Node], new_index: Optional[int] = None - ) -> None: - """Insert under parent node.""" - new_index = new_index or self.size - if parent_node is None: - self.root_nodes[new_index] = node.get_doc_id() - self.node_id_to_children_ids[node.get_doc_id()] = [] - else: - if parent_node.doc_id not in self.node_id_to_children_ids: - self.node_id_to_children_ids[parent_node.get_doc_id()] = [] - self.node_id_to_children_ids[parent_node.get_doc_id()].append( - node.get_doc_id() - ) - - self.all_nodes[new_index] = node.get_doc_id() - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.TREE - - -@dataclass -class KeywordTable(V2IndexStruct): - """A table of keywords mapping keywords to text chunks.""" - - table: Dict[str, Set[str]] = field(default_factory=dict) - - def add_node(self, keywords: List[str], node: Node) -> None: - """Add text to table.""" - for keyword in keywords: - if keyword not in self.table: - self.table[keyword] = set() - self.table[keyword].add(node.get_doc_id()) - - @property - def node_ids(self) -> Set[str]: - """Get all node ids.""" - return set.union(*self.table.values()) - - @property - def keywords(self) -> Set[str]: - """Get all keywords in the table.""" - return set(self.table.keys()) - - @property - def size(self) -> int: - """Get the size of the table.""" - return len(self.table) - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.KEYWORD_TABLE - - -@dataclass -class IndexList(V2IndexStruct): - """A list of documents.""" - - nodes: List[str] = field(default_factory=list) - - def add_node(self, node: Node) -> None: - """Add text to table, return current position in list.""" - # don't worry about child indices for now, nodes are all in order - self.nodes.append(node.get_doc_id()) - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.LIST - - -@dataclass -class IndexDict(V2IndexStruct): - """A simple dictionary of documents.""" - - # mapping from vector store id to node id - nodes_dict: Dict[str, str] = field(default_factory=dict) - # mapping from doc_id to vector store id - doc_id_dict: Dict[str, List[str]] = field(default_factory=dict) - - # TODO: temporary hack to store embeddings for simple vector index - # this should be empty for all other indices - embeddings_dict: Dict[str, List[float]] = field(default_factory=dict) - - def add_node( - self, - node: Node, - text_id: Optional[str] = None, - ) -> str: - """Add text to table, return current position in list.""" - # # don't worry about child indices for now, nodes are all in order - # self.nodes_dict[int_id] = node - vector_id = text_id if text_id is not None else node.get_doc_id() - self.nodes_dict[vector_id] = node.get_doc_id() - if node.ref_doc_id is not None: - if node.ref_doc_id not in self.doc_id_dict: - self.doc_id_dict[node.ref_doc_id] = [] - self.doc_id_dict[node.ref_doc_id].append(vector_id) - - return vector_id - - def delete(self, doc_id: str) -> None: - """Delete a Node.""" - if doc_id not in self.doc_id_dict: - return - for vector_id in self.doc_id_dict[doc_id]: - del self.nodes_dict[vector_id] - del self.doc_id_dict[doc_id] - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.VECTOR_STORE - - -@dataclass -class KG(V2IndexStruct): - """A table of keywords mapping keywords to text chunks.""" - - # Unidirectional - - table: Dict[str, Set[str]] = field(default_factory=dict) - # text_chunks: Dict[str, Node] = field(default_factory=dict) - rel_map: Dict[str, List[Tuple[str, str]]] = field(default_factory=dict) - embedding_dict: Dict[str, List[float]] = field(default_factory=dict) - - @property - def node_ids(self) -> Set[str]: - """Get all node ids.""" - return set.union(*self.table.values()) - - def add_to_embedding_dict(self, triplet_str: str, embedding: List[float]) -> None: - """Add embedding to dict.""" - self.embedding_dict[triplet_str] = embedding - - def upsert_triplet(self, triplet: Tuple[str, str, str]) -> None: - """Upsert a knowledge triplet to the graph.""" - subj, relationship, obj = triplet - if subj not in self.rel_map: - self.rel_map[subj] = [] - self.rel_map[subj].append((obj, relationship)) - - def add_node(self, keywords: List[str], node: Node) -> None: - """Add text to table.""" - node_id = node.get_doc_id() - for keyword in keywords: - if keyword not in self.table: - self.table[keyword] = set() - self.table[keyword].add(node_id) - # self.text_chunks[node_id] = node - - def get_rel_map_texts(self, keyword: str) -> List[str]: - """Get the corresponding knowledge for a given keyword.""" - # NOTE: return a single node for now - if keyword not in self.rel_map: - return [] - texts = [] - for obj, rel in self.rel_map[keyword]: - texts.append(str((keyword, rel, obj))) - return texts - - def get_rel_map_tuples(self, keyword: str) -> List[Tuple[str, str]]: - """Get the corresponding knowledge for a given keyword.""" - # NOTE: return a single node for now - if keyword not in self.rel_map: - return [] - return self.rel_map[keyword] - - def get_node_ids(self, keyword: str, depth: int = 1) -> List[str]: - """Get the corresponding knowledge for a given keyword.""" - if depth > 1: - raise ValueError("Depth > 1 not supported yet.") - if keyword not in self.table: - return [] - keywords = [keyword] - # some keywords may correspond to a leaf node, may not be in rel_map - if keyword in self.rel_map: - keywords.extend([child for child, _ in self.rel_map[keyword]]) - - node_ids: List[str] = [] - for keyword in keywords: - for node_id in self.table.get(keyword, set()): - node_ids.append(node_id) - # TODO: Traverse (with depth > 1) - return node_ids - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.KG - - -@dataclass -class EmptyIndex(IndexDict): - """Empty index.""" - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.EMPTY diff --git a/llama_index/data_structs/node_v2.py b/llama_index/data_structs/node.py similarity index 100% rename from llama_index/data_structs/node_v2.py rename to llama_index/data_structs/node.py diff --git a/llama_index/data_structs/registry.py b/llama_index/data_structs/registry.py index 0c4d8c91c709c..a8ae5e0a764fb 100644 --- a/llama_index/data_structs/registry.py +++ b/llama_index/data_structs/registry.py @@ -2,20 +2,20 @@ from typing import Dict, Type -from llama_index.data_structs.data_structs_v2 import ( +from llama_index.data_structs.data_structs import ( KG, EmptyIndex, IndexDict, IndexGraph, IndexList, KeywordTable, - V2IndexStruct, + IndexStruct, ) from llama_index.data_structs.struct_type import IndexStructType -from llama_index.data_structs.table_v2 import PandasStructTable, SQLStructTable +from llama_index.data_structs.table import PandasStructTable, SQLStructTable -INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS: Dict[IndexStructType, Type[V2IndexStruct]] = { +INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS: Dict[IndexStructType, Type[IndexStruct]] = { IndexStructType.TREE: IndexGraph, IndexStructType.LIST: IndexList, IndexStructType.KEYWORD_TABLE: KeywordTable, diff --git a/llama_index/data_structs/table.py b/llama_index/data_structs/table.py index dbb7155f61067..f29d5212c77dd 100644 --- a/llama_index/data_structs/table.py +++ b/llama_index/data_structs/table.py @@ -6,6 +6,7 @@ from dataclasses_json import DataClassJsonMixin from llama_index.data_structs.data_structs import IndexStruct +from llama_index.data_structs.struct_type import IndexStructType @dataclass @@ -28,10 +29,10 @@ class SQLStructTable(BaseStructTable): context_dict: Dict[str, str] = field(default_factory=dict) @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" # TODO: consolidate with IndexStructType - return "sql" + return IndexStructType.SQL @dataclass @@ -39,6 +40,6 @@ class PandasStructTable(BaseStructTable): """Pandas struct outputs.""" @classmethod - def get_type(cls) -> str: + def get_type(cls) -> IndexStructType: """Get type.""" - return "pandas" + return IndexStructType.PANDAS diff --git a/llama_index/data_structs/table_v2.py b/llama_index/data_structs/table_v2.py deleted file mode 100644 index bfcf600a489da..0000000000000 --- a/llama_index/data_structs/table_v2.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Struct store schema.""" - -from dataclasses import dataclass, field -from typing import Any, Dict - -from dataclasses_json import DataClassJsonMixin - -from llama_index.data_structs.data_structs_v2 import V2IndexStruct -from llama_index.data_structs.struct_type import IndexStructType - - -@dataclass -class StructDatapoint(DataClassJsonMixin): - """Struct outputs.""" - - # map from field name to StructValue - fields: Dict[str, Any] - - -@dataclass -class BaseStructTable(V2IndexStruct): - """Struct outputs.""" - - -@dataclass -class SQLStructTable(BaseStructTable): - """SQL struct outputs.""" - - context_dict: Dict[str, str] = field(default_factory=dict) - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - # TODO: consolidate with IndexStructType - return IndexStructType.SQL - - -@dataclass -class PandasStructTable(BaseStructTable): - """Pandas struct outputs.""" - - @classmethod - def get_type(cls) -> IndexStructType: - """Get type.""" - return IndexStructType.PANDAS diff --git a/llama_index/evaluation/dataset_generation.py b/llama_index/evaluation/dataset_generation.py index df39a86d78f40..94b087f8acd91 100644 --- a/llama_index/evaluation/dataset_generation.py +++ b/llama_index/evaluation/dataset_generation.py @@ -12,7 +12,7 @@ LLMPredictor, ) from llama_index.indices.postprocessor.node import KeywordNodePostprocessor -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from langchain.chat_models import ChatOpenAI diff --git a/llama_index/experimental/evaporate/base.py b/llama_index/experimental/evaporate/base.py index a9d4551b264a7..de341d9f8fc62 100644 --- a/llama_index/experimental/evaporate/base.py +++ b/llama_index/experimental/evaporate/base.py @@ -4,7 +4,7 @@ from llama_index.prompts.prompts import ( QuestionAnswerPrompt, ) -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from typing import Optional, List, Dict, Tuple from llama_index.experimental.evaporate.prompts import ( SchemaIDPrompt, diff --git a/llama_index/indices/base.py b/llama_index/indices/base.py index da5c376f6f7a6..c8eea3309bff3 100644 --- a/llama_index/indices/base.py +++ b/llama_index/indices/base.py @@ -4,8 +4,8 @@ from typing import Any, Generic, List, Optional, Sequence, Type, TypeVar from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.data_structs_v2 import V2IndexStruct -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexStruct +from llama_index.data_structs.node import Node from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.base import BaseQueryEngine from llama_index.indices.service_context import ServiceContext @@ -14,7 +14,7 @@ from llama_index.storage.storage_context import StorageContext from llama_index.token_counter.token_counter import llm_token_counter -IS = TypeVar("IS", bound=V2IndexStruct) +IS = TypeVar("IS", bound=IndexStruct) IndexType = TypeVar("IndexType", bound="BaseGPTIndex") logger = logging.getLogger(__name__) diff --git a/llama_index/indices/base_retriever.py b/llama_index/indices/base_retriever.py index cfc554fd4037b..c195e5bd84e24 100644 --- a/llama_index/indices/base_retriever.py +++ b/llama_index/indices/base_retriever.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from typing import List -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.query.schema import QueryBundle, QueryType diff --git a/llama_index/indices/common/struct_store/base.py b/llama_index/indices/common/struct_store/base.py index 52d56bfff308f..0a1095ec453d7 100644 --- a/llama_index/indices/common/struct_store/base.py +++ b/llama_index/indices/common/struct_store/base.py @@ -5,7 +5,7 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, cast from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.data_structs.table import StructDatapoint from llama_index.indices.response.response_builder import get_response_builder from llama_index.indices.service_context import ServiceContext diff --git a/llama_index/indices/common_tree/base.py b/llama_index/indices/common_tree/base.py index e43bbcfb0b4ed..ffd30c6fd167d 100644 --- a/llama_index/indices/common_tree/base.py +++ b/llama_index/indices/common_tree/base.py @@ -7,8 +7,8 @@ from llama_index.async_utils import run_async_tasks from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.data_structs_v2 import IndexGraph -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexGraph +from llama_index.data_structs.node import Node from llama_index.storage.docstore import BaseDocumentStore from llama_index.storage.docstore.registry import get_default_docstore from llama_index.indices.service_context import ServiceContext diff --git a/llama_index/indices/composability/graph.py b/llama_index/indices/composability/graph.py index a9ee763cada45..62ce98d9c7ca9 100644 --- a/llama_index/indices/composability/graph.py +++ b/llama_index/indices/composability/graph.py @@ -2,8 +2,8 @@ from typing import Any, Dict, List, Optional, Sequence, Type, cast -from llama_index.data_structs.data_structs_v2 import V2IndexStruct -from llama_index.data_structs.node_v2 import IndexNode, DocumentRelationship +from llama_index.data_structs.data_structs import IndexStruct +from llama_index.data_structs.node import IndexNode, DocumentRelationship from llama_index.indices.base import BaseGPTIndex from llama_index.indices.query.base import BaseQueryEngine from llama_index.indices.service_context import ServiceContext @@ -34,7 +34,7 @@ def root_index(self) -> BaseGPTIndex: return self._all_indices[self._root_id] @property - def index_struct(self) -> V2IndexStruct: + def index_struct(self) -> IndexStruct: return self._all_indices[self._root_id].index_struct @property @@ -73,7 +73,7 @@ def from_indices( # construct index nodes index_nodes = [] for index, summary in zip(children_indices, index_summaries): - assert isinstance(index.index_struct, V2IndexStruct) + assert isinstance(index.index_struct, IndexStruct) index_node = IndexNode( text=summary, index_id=index.index_id, diff --git a/llama_index/indices/empty/base.py b/llama_index/indices/empty/base.py index b75542def35f6..ac29a85f13939 100644 --- a/llama_index/indices/empty/base.py +++ b/llama_index/indices/empty/base.py @@ -7,8 +7,8 @@ from typing import Any, Optional, Sequence -from llama_index.data_structs.data_structs_v2 import EmptyIndex -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import EmptyIndex +from llama_index.data_structs.node import Node from llama_index.indices.base import BaseGPTIndex from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.service_context import ServiceContext diff --git a/llama_index/indices/empty/retrievers.py b/llama_index/indices/empty/retrievers.py index 09b0552a1db83..b2713e7e6d81e 100644 --- a/llama_index/indices/empty/retrievers.py +++ b/llama_index/indices/empty/retrievers.py @@ -2,7 +2,7 @@ from typing import Any, List, Optional from llama_index.indices.base_retriever import BaseRetriever -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.query.schema import QueryBundle from llama_index.prompts.default_prompts import DEFAULT_SIMPLE_INPUT_PROMPT from llama_index.prompts.prompts import SimpleInputPrompt diff --git a/llama_index/indices/keyword_table/base.py b/llama_index/indices/keyword_table/base.py index a105bdaaeb959..bf1d034926d9a 100644 --- a/llama_index/indices/keyword_table/base.py +++ b/llama_index/indices/keyword_table/base.py @@ -14,8 +14,8 @@ from llama_index.async_utils import run_async_tasks from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.data_structs_v2 import KeywordTable -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import KeywordTable +from llama_index.data_structs.node import Node from llama_index.indices.base import BaseGPTIndex from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.keyword_table.utils import extract_keywords_given_response diff --git a/llama_index/indices/keyword_table/retrievers.py b/llama_index/indices/keyword_table/retrievers.py index 9ffeb9d505fe4..56f5f2752eae2 100644 --- a/llama_index/indices/keyword_table/retrievers.py +++ b/llama_index/indices/keyword_table/retrievers.py @@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.keyword_table.base import ( BaseGPTKeywordTableIndex, diff --git a/llama_index/indices/knowledge_graph/base.py b/llama_index/indices/knowledge_graph/base.py index 7f5db52432cc1..e2107e9acb203 100644 --- a/llama_index/indices/knowledge_graph/base.py +++ b/llama_index/indices/knowledge_graph/base.py @@ -12,8 +12,8 @@ from typing import Any, List, Optional, Sequence, Tuple from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.data_structs_v2 import KG -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import KG +from llama_index.data_structs.node import Node from llama_index.indices.base import BaseGPTIndex from llama_index.indices.base_retriever import BaseRetriever from llama_index.prompts.default_prompts import ( diff --git a/llama_index/indices/knowledge_graph/retrievers.py b/llama_index/indices/knowledge_graph/retrievers.py index 99a970e6795ad..eb8a2ab53134c 100644 --- a/llama_index/indices/knowledge_graph/retrievers.py +++ b/llama_index/indices/knowledge_graph/retrievers.py @@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.keyword_table.utils import extract_keywords_given_response from llama_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex diff --git a/llama_index/indices/list/base.py b/llama_index/indices/list/base.py index a9c092d65e5dc..53d4face35f9f 100644 --- a/llama_index/indices/list/base.py +++ b/llama_index/indices/list/base.py @@ -8,8 +8,8 @@ from enum import Enum from typing import Any, Optional, Sequence, Union -from llama_index.data_structs.data_structs_v2 import IndexList -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexList +from llama_index.data_structs.node import Node from llama_index.indices.base import BaseGPTIndex from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.service_context import ServiceContext diff --git a/llama_index/indices/list/retrievers.py b/llama_index/indices/list/retrievers.py index 86d3f05117ee2..85a6a6352adfa 100644 --- a/llama_index/indices/list/retrievers.py +++ b/llama_index/indices/list/retrievers.py @@ -3,7 +3,7 @@ from typing import Any, List, Optional, Tuple from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.embedding_utils import ( get_top_k_embeddings, diff --git a/llama_index/indices/postprocessor/node.py b/llama_index/indices/postprocessor/node.py index 630c0ebf90dc0..c4b678cebba7b 100644 --- a/llama_index/indices/postprocessor/node.py +++ b/llama_index/indices/postprocessor/node.py @@ -12,7 +12,7 @@ from llama_index.indices.service_context import ServiceContext from llama_index.prompts.prompts import QuestionAnswerPrompt, RefinePrompt from llama_index.storage.docstore import BaseDocumentStore -from llama_index.data_structs.node_v2 import DocumentRelationship, NodeWithScore +from llama_index.data_structs.node import DocumentRelationship, NodeWithScore from llama_index.indices.postprocessor.base import BasePostprocessor from llama_index.indices.response.response_builder import get_response_builder diff --git a/llama_index/indices/postprocessor/node_recency.py b/llama_index/indices/postprocessor/node_recency.py index ae1a9d7ee3e11..25f0163cc14ad 100644 --- a/llama_index/indices/postprocessor/node_recency.py +++ b/llama_index/indices/postprocessor/node_recency.py @@ -2,7 +2,7 @@ from llama_index.indices.postprocessor.node import BaseNodePostprocessor from llama_index.indices.service_context import ServiceContext -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from pydantic import Field from typing import Optional, Dict, List, Set import pandas as pd diff --git a/llama_index/indices/postprocessor/pii.py b/llama_index/indices/postprocessor/pii.py index cd8af849e3658..c6c274ad1da91 100644 --- a/llama_index/indices/postprocessor/pii.py +++ b/llama_index/indices/postprocessor/pii.py @@ -1,7 +1,7 @@ """PII postprocessor.""" from llama_index.indices.postprocessor.node import BaseNodePostprocessor -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from typing import List, Optional, Dict, Tuple, Callable from llama_index.indices.service_context import ServiceContext from llama_index.prompts.prompts import QuestionAnswerPrompt diff --git a/llama_index/indices/prompt_helper.py b/llama_index/indices/prompt_helper.py index 164685a85bc00..65ee1d1d2e9fb 100644 --- a/llama_index/indices/prompt_helper.py +++ b/llama_index/indices/prompt_helper.py @@ -8,7 +8,7 @@ from typing import Callable, List, Optional, Sequence from llama_index.constants import MAX_CHUNK_OVERLAP -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.langchain_helpers.chain_wrapper import LLMPredictor from llama_index.langchain_helpers.text_splitter import TokenTextSplitter from llama_index.prompts.base import Prompt diff --git a/llama_index/indices/query/base.py b/llama_index/indices/query/base.py index 3d5c5748f70dd..982a8f996c5eb 100644 --- a/llama_index/indices/query/base.py +++ b/llama_index/indices/query/base.py @@ -3,7 +3,7 @@ import logging from abc import ABC, abstractmethod from typing import List, Optional, Sequence -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.query.schema import QueryBundle, QueryType from llama_index.response.schema import ( diff --git a/llama_index/indices/query/response_synthesis.py b/llama_index/indices/query/response_synthesis.py index 79a9c5b9266a1..58922835f6ed3 100644 --- a/llama_index/indices/query/response_synthesis.py +++ b/llama_index/indices/query/response_synthesis.py @@ -8,7 +8,7 @@ Sequence, ) -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from llama_index.indices.postprocessor.node import BaseNodePostprocessor from llama_index.indices.query.schema import QueryBundle from llama_index.indices.response.response_builder import ( diff --git a/llama_index/indices/registry.py b/llama_index/indices/registry.py index 5207d6f42aa92..d08fdfa0e6c62 100644 --- a/llama_index/indices/registry.py +++ b/llama_index/indices/registry.py @@ -2,17 +2,17 @@ from typing import Dict, Type -from llama_index.data_structs.data_structs_v2 import ( +from llama_index.data_structs.data_structs import ( KG, EmptyIndex, IndexDict, IndexGraph, IndexList, KeywordTable, - V2IndexStruct, + IndexStruct, ) from llama_index.data_structs.struct_type import IndexStructType -from llama_index.data_structs.table_v2 import PandasStructTable, SQLStructTable +from llama_index.data_structs.table import PandasStructTable, SQLStructTable from llama_index.indices.base import BaseGPTIndex from llama_index.indices.empty.base import GPTEmptyIndex from llama_index.indices.keyword_table.base import GPTKeywordTableIndex @@ -23,7 +23,7 @@ from llama_index.indices.tree.base import GPTTreeIndex from llama_index.indices.vector_store.base import GPTVectorStoreIndex -INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS: Dict[IndexStructType, Type[V2IndexStruct]] = { +INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS: Dict[IndexStructType, Type[IndexStruct]] = { IndexStructType.TREE: IndexGraph, IndexStructType.LIST: IndexList, IndexStructType.KEYWORD_TABLE: KeywordTable, diff --git a/llama_index/indices/response/response_builder.py b/llama_index/indices/response/response_builder.py index 0d422dbbf9a25..ef00a18590207 100644 --- a/llama_index/indices/response/response_builder.py +++ b/llama_index/indices/response/response_builder.py @@ -12,8 +12,8 @@ from typing import Any, Dict, Generator, List, Optional, Sequence, Tuple, cast from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.data_structs_v2 import IndexGraph -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexGraph +from llama_index.data_structs.node import Node from llama_index.storage.docstore.registry import get_default_docstore from llama_index.indices.common_tree.base import GPTTreeIndexBuilder from llama_index.indices.response.type import ResponseMode diff --git a/llama_index/indices/struct_store/base.py b/llama_index/indices/struct_store/base.py index 6183ad9a1ce29..fff1e84971721 100644 --- a/llama_index/indices/struct_store/base.py +++ b/llama_index/indices/struct_store/base.py @@ -3,8 +3,8 @@ import re from typing import Any, Callable, Dict, Generic, Optional, Sequence, TypeVar -from llama_index.data_structs.node_v2 import Node -from llama_index.data_structs.table_v2 import BaseStructTable +from llama_index.data_structs.node import Node +from llama_index.data_structs.table import BaseStructTable from llama_index.indices.base import BaseGPTIndex from llama_index.indices.service_context import ServiceContext from llama_index.prompts.default_prompts import DEFAULT_SCHEMA_EXTRACT_PROMPT diff --git a/llama_index/indices/struct_store/pandas.py b/llama_index/indices/struct_store/pandas.py index ddd91c6e01e0c..82495c3f71432 100644 --- a/llama_index/indices/struct_store/pandas.py +++ b/llama_index/indices/struct_store/pandas.py @@ -2,8 +2,8 @@ from typing import Any, Optional, Sequence -from llama_index.data_structs.node_v2 import Node -from llama_index.data_structs.table_v2 import PandasStructTable +from llama_index.data_structs.node import Node +from llama_index.data_structs.table import PandasStructTable from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.base import BaseQueryEngine from llama_index.indices.struct_store.base import BaseGPTStructStoreIndex diff --git a/llama_index/indices/struct_store/sql.py b/llama_index/indices/struct_store/sql.py index fb7313ca77f25..c6387849085c1 100644 --- a/llama_index/indices/struct_store/sql.py +++ b/llama_index/indices/struct_store/sql.py @@ -3,8 +3,8 @@ from typing import Any, Optional, Sequence, Union from collections import defaultdict -from llama_index.data_structs.node_v2 import Node -from llama_index.data_structs.table_v2 import SQLStructTable +from llama_index.data_structs.node import Node +from llama_index.data_structs.table import SQLStructTable from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.common.struct_store.schema import SQLContextContainer from llama_index.indices.common.struct_store.sql import SQLStructDatapointExtractor diff --git a/llama_index/indices/tree/all_leaf_retriever.py b/llama_index/indices/tree/all_leaf_retriever.py index 80dcad2b4a050..703976d5e65cb 100644 --- a/llama_index/indices/tree/all_leaf_retriever.py +++ b/llama_index/indices/tree/all_leaf_retriever.py @@ -4,8 +4,8 @@ from typing import Any, List, cast -from llama_index.data_structs.data_structs_v2 import IndexGraph -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.data_structs import IndexGraph +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.schema import QueryBundle from llama_index.indices.utils import get_sorted_node_list diff --git a/llama_index/indices/tree/base.py b/llama_index/indices/tree/base.py index 8c80a019473f6..831e4696c91b6 100644 --- a/llama_index/indices/tree/base.py +++ b/llama_index/indices/tree/base.py @@ -4,8 +4,8 @@ from typing import Any, Optional, Sequence, Union # from llama_index.data_structs.data_structs import IndexGraph -from llama_index.data_structs.data_structs_v2 import IndexGraph -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexGraph +from llama_index.data_structs.node import Node from llama_index.indices.base import BaseGPTIndex from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.common_tree.base import GPTTreeIndexBuilder diff --git a/llama_index/indices/tree/inserter.py b/llama_index/indices/tree/inserter.py index d44003ba80ecd..06cde714380af 100644 --- a/llama_index/indices/tree/inserter.py +++ b/llama_index/indices/tree/inserter.py @@ -2,8 +2,8 @@ from typing import Optional, Sequence -from llama_index.data_structs.data_structs_v2 import IndexGraph -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexGraph +from llama_index.data_structs.node import Node from llama_index.storage.docstore import BaseDocumentStore from llama_index.storage.docstore.registry import get_default_docstore from llama_index.indices.service_context import ServiceContext diff --git a/llama_index/indices/tree/select_leaf_embedding_retriever.py b/llama_index/indices/tree/select_leaf_embedding_retriever.py index 1316fbbf6bc74..5fbd3b63a346a 100644 --- a/llama_index/indices/tree/select_leaf_embedding_retriever.py +++ b/llama_index/indices/tree/select_leaf_embedding_retriever.py @@ -4,7 +4,7 @@ from typing import Dict, List, Tuple, cast -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.indices.query.schema import QueryBundle from llama_index.indices.tree.select_leaf_retriever import ( TreeSelectLeafRetriever, diff --git a/llama_index/indices/tree/select_leaf_retriever.py b/llama_index/indices/tree/select_leaf_retriever.py index 3ae95de36250e..f7043fa99bae9 100644 --- a/llama_index/indices/tree/select_leaf_retriever.py +++ b/llama_index/indices/tree/select_leaf_retriever.py @@ -5,7 +5,7 @@ from langchain.input import print_text -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.schema import QueryBundle from llama_index.indices.response.response_builder import get_response_builder diff --git a/llama_index/indices/tree/tree_root_retriever.py b/llama_index/indices/tree/tree_root_retriever.py index 770b7cd4e0757..39006d9aee16f 100644 --- a/llama_index/indices/tree/tree_root_retriever.py +++ b/llama_index/indices/tree/tree_root_retriever.py @@ -2,7 +2,7 @@ import logging from typing import Any, List -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.schema import QueryBundle from llama_index.indices.utils import get_sorted_node_list diff --git a/llama_index/indices/utils.py b/llama_index/indices/utils.py index 9f1e237810e7c..6435485252d84 100644 --- a/llama_index/indices/utils.py +++ b/llama_index/indices/utils.py @@ -3,7 +3,7 @@ import re from typing import Dict, List, Optional, Set -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.utils import globals_helper, truncate_text from llama_index.vector_stores.types import VectorStoreQueryResult diff --git a/llama_index/indices/vector_store/base.py b/llama_index/indices/vector_store/base.py index 0a9b3af9a8b4b..b1bc332a59acf 100644 --- a/llama_index/indices/vector_store/base.py +++ b/llama_index/indices/vector_store/base.py @@ -8,8 +8,8 @@ from llama_index.callbacks.schema import CBEventType from llama_index.async_utils import run_async_tasks -from llama_index.data_structs.data_structs_v2 import IndexDict -from llama_index.data_structs.node_v2 import ImageNode, IndexNode, Node +from llama_index.data_structs.data_structs import IndexDict +from llama_index.data_structs.node import ImageNode, IndexNode, Node from llama_index.indices.base import BaseGPTIndex from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.service_context import ServiceContext diff --git a/llama_index/indices/vector_store/retrievers.py b/llama_index/indices/vector_store/retrievers.py index 082fd5e0fbb2c..e4c4e8a1f1805 100644 --- a/llama_index/indices/vector_store/retrievers.py +++ b/llama_index/indices/vector_store/retrievers.py @@ -5,8 +5,8 @@ from llama_index.constants import DEFAULT_SIMILARITY_TOP_K from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.data_structs_v2 import IndexDict -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.data_structs import IndexDict +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.schema import QueryBundle from llama_index.indices.utils import log_vector_store_query_result diff --git a/llama_index/node_parser/interface.py b/llama_index/node_parser/interface.py index af0eb841d3f73..bb0a4337a64b6 100644 --- a/llama_index/node_parser/interface.py +++ b/llama_index/node_parser/interface.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.readers.schema.base import Document diff --git a/llama_index/node_parser/node_utils.py b/llama_index/node_parser/node_utils.py index 7c4f8b2b6eb46..d75642ee21f0b 100644 --- a/llama_index/node_parser/node_utils.py +++ b/llama_index/node_parser/node_utils.py @@ -4,7 +4,7 @@ import logging from typing import List -from llama_index.data_structs.node_v2 import DocumentRelationship, ImageNode, Node +from llama_index.data_structs.node import DocumentRelationship, ImageNode, Node from llama_index.langchain_helpers.text_splitter import ( TextSplit, TextSplitter, diff --git a/llama_index/node_parser/simple.py b/llama_index/node_parser/simple.py index b2e2738416700..bc33f6251a56a 100644 --- a/llama_index/node_parser/simple.py +++ b/llama_index/node_parser/simple.py @@ -1,7 +1,7 @@ """Simple node parser.""" from typing import List, Optional, Sequence -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.langchain_helpers.text_splitter import TextSplitter, TokenTextSplitter from llama_index.node_parser.node_utils import get_nodes_from_document from llama_index.readers.schema.base import Document diff --git a/llama_index/query_engine/graph_query_engine.py b/llama_index/query_engine/graph_query_engine.py index 0e2b12e71a3d5..5a954bb3fd785 100644 --- a/llama_index/query_engine/graph_query_engine.py +++ b/llama_index/query_engine/graph_query_engine.py @@ -1,6 +1,6 @@ from typing import Dict, List, Optional, Tuple -from llama_index.data_structs.node_v2 import IndexNode, Node, NodeWithScore +from llama_index.data_structs.node import IndexNode, Node, NodeWithScore from llama_index.indices.composability.graph import ComposableGraph from llama_index.indices.query.base import BaseQueryEngine from llama_index.indices.query.schema import QueryBundle diff --git a/llama_index/query_engine/multistep_query_engine.py b/llama_index/query_engine/multistep_query_engine.py index 76855621e9ca3..da89b30c541b2 100644 --- a/llama_index/query_engine/multistep_query_engine.py +++ b/llama_index/query_engine/multistep_query_engine.py @@ -1,5 +1,5 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, cast -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from llama_index.indices.query.base import BaseQueryEngine from llama_index.indices.query.query_transform.base import StepDecomposeQueryTransform from llama_index.indices.query.schema import QueryBundle diff --git a/llama_index/query_engine/retriever_query_engine.py b/llama_index/query_engine/retriever_query_engine.py index 4f671997d8949..d7a9ba5c2ba8f 100644 --- a/llama_index/query_engine/retriever_query_engine.py +++ b/llama_index/query_engine/retriever_query_engine.py @@ -2,7 +2,7 @@ from llama_index.callbacks.base import CallbackManager from llama_index.callbacks.schema import CBEventType -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.postprocessor.node import BaseNodePostprocessor from llama_index.indices.query.base import BaseQueryEngine diff --git a/llama_index/query_engine/transform_query_engine.py b/llama_index/query_engine/transform_query_engine.py index 67c78c6c0df8d..a048fd52ef2c9 100644 --- a/llama_index/query_engine/transform_query_engine.py +++ b/llama_index/query_engine/transform_query_engine.py @@ -1,5 +1,5 @@ from typing import List, Optional, Sequence -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.query.base import BaseQueryEngine from llama_index.indices.query.query_transform.base import BaseQueryTransform from llama_index.indices.query.schema import QueryBundle diff --git a/llama_index/readers/make_com/wrapper.py b/llama_index/readers/make_com/wrapper.py index 126e7146ec6de..746007cc702eb 100644 --- a/llama_index/readers/make_com/wrapper.py +++ b/llama_index/readers/make_com/wrapper.py @@ -7,7 +7,7 @@ from typing import Any, List, Optional import requests -from llama_index.data_structs.node_v2 import Node, NodeWithScore +from llama_index.data_structs.node import Node, NodeWithScore from llama_index.readers.base import BaseReader from llama_index.readers.schema.base import Document diff --git a/llama_index/readers/weaviate/client.py b/llama_index/readers/weaviate/client.py index 415943dd060df..7dc43c4f61357 100644 --- a/llama_index/readers/weaviate/client.py +++ b/llama_index/readers/weaviate/client.py @@ -8,8 +8,8 @@ from dataclasses import field from typing import Any, Dict, List, Optional, cast -from llama_index.data_structs.data_structs_v2 import Node -from llama_index.data_structs.node_v2 import DocumentRelationship +from llama_index.data_structs.data_structs import Node +from llama_index.data_structs.node import DocumentRelationship from llama_index.readers.weaviate.utils import ( get_by_id, parse_get_response, diff --git a/llama_index/response/notebook_utils.py b/llama_index/response/notebook_utils.py index 49da84e9f1897..5996cd9641ec6 100644 --- a/llama_index/response/notebook_utils.py +++ b/llama_index/response/notebook_utils.py @@ -2,7 +2,7 @@ from typing import Any, Dict, Tuple from IPython.display import Markdown, display -from llama_index.data_structs.node_v2 import ImageNode, NodeWithScore +from llama_index.data_structs.node import ImageNode, NodeWithScore from llama_index.img_utils import b64_2_img from llama_index.response.schema import Response diff --git a/llama_index/response/schema.py b/llama_index/response/schema.py index a325470857f39..103df1e6df518 100644 --- a/llama_index/response/schema.py +++ b/llama_index/response/schema.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, Generator, List, Optional, Union -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.utils import truncate_text diff --git a/llama_index/retrievers/transform_retriever.py b/llama_index/retrievers/transform_retriever.py index 8c4ce5c4cee99..a8a49642445d2 100644 --- a/llama_index/retrievers/transform_retriever.py +++ b/llama_index/retrievers/transform_retriever.py @@ -1,5 +1,5 @@ from typing import List, Optional -from llama_index.data_structs.node_v2 import NodeWithScore +from llama_index.data_structs.node import NodeWithScore from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.query.query_transform.base import BaseQueryTransform from llama_index.indices.query.schema import QueryBundle diff --git a/llama_index/storage/docstore/types.py b/llama_index/storage/docstore/types.py index a43842e49783a..bedb98e93e36a 100644 --- a/llama_index/storage/docstore/types.py +++ b/llama_index/storage/docstore/types.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from typing import Dict, List, Optional, Sequence -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.schema import BaseDocument import os diff --git a/llama_index/storage/docstore/utils.py b/llama_index/storage/docstore/utils.py index ad430cb6366de..457e5b22b7d53 100644 --- a/llama_index/storage/docstore/utils.py +++ b/llama_index/storage/docstore/utils.py @@ -1,5 +1,5 @@ from llama_index.constants import DATA_KEY, TYPE_KEY -from llama_index.data_structs.node_v2 import ImageNode, IndexNode, Node +from llama_index.data_structs.node import ImageNode, IndexNode, Node from llama_index.readers.schema.base import Document from llama_index.schema import BaseDocument diff --git a/llama_index/storage/index_store/keyval_index_store.py b/llama_index/storage/index_store/keyval_index_store.py index 714f8ed384b64..2a5616100d3b5 100644 --- a/llama_index/storage/index_store/keyval_index_store.py +++ b/llama_index/storage/index_store/keyval_index_store.py @@ -1,5 +1,5 @@ from typing import List, Optional -from llama_index.data_structs.data_structs_v2 import V2IndexStruct +from llama_index.data_structs.data_structs import IndexStruct from llama_index.storage.index_store.types import BaseIndexStore from llama_index.storage.index_store.utils import ( index_struct_to_json, @@ -25,11 +25,11 @@ def __init__(self, kvstore: BaseKVStore, namespace: Optional[str] = None) -> Non namespace = namespace or DEFAULT_NAMESPACE self._collection = f"{namespace}/data" - def add_index_struct(self, index_struct: V2IndexStruct) -> None: + def add_index_struct(self, index_struct: IndexStruct) -> None: """Add an index struct. Args: - index_struct (V2IndexStruct): index struct + index_struct (IndexStruct): index struct """ key = index_struct.index_id @@ -47,7 +47,7 @@ def delete_index_struct(self, key: str) -> None: def get_index_struct( self, struct_id: Optional[str] = None - ) -> Optional[V2IndexStruct]: + ) -> Optional[IndexStruct]: """Get an index struct. Args: @@ -64,11 +64,11 @@ def get_index_struct( return None return json_to_index_struct(json) - def index_structs(self) -> List[V2IndexStruct]: + def index_structs(self) -> List[IndexStruct]: """Get all index structs. Returns: - List[V2IndexStruct]: index structs + List[IndexStruct]: index structs """ jsons = self._kvstore.get_all(collection=self._collection) diff --git a/llama_index/storage/index_store/types.py b/llama_index/storage/index_store/types.py index 24d9f0727e2ff..1ceca56132dcb 100644 --- a/llama_index/storage/index_store/types.py +++ b/llama_index/storage/index_store/types.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import List, Optional -from llama_index.data_structs.data_structs_v2 import V2IndexStruct +from llama_index.data_structs.data_structs import IndexStruct import os DEFAULT_PERSIST_DIR = "./storage" @@ -11,11 +11,11 @@ class BaseIndexStore(ABC): @abstractmethod - def index_structs(self) -> List[V2IndexStruct]: + def index_structs(self) -> List[IndexStruct]: pass @abstractmethod - def add_index_struct(self, index_struct: V2IndexStruct) -> None: + def add_index_struct(self, index_struct: IndexStruct) -> None: pass @abstractmethod @@ -25,7 +25,7 @@ def delete_index_struct(self, key: str) -> None: @abstractmethod def get_index_struct( self, struct_id: Optional[str] = None - ) -> Optional[V2IndexStruct]: + ) -> Optional[IndexStruct]: pass def persist(self, persist_path: str = DEFAULT_PERSIST_PATH) -> None: diff --git a/llama_index/storage/index_store/utils.py b/llama_index/storage/index_store/utils.py index 706bcd91c5259..45455b09a1a29 100644 --- a/llama_index/storage/index_store/utils.py +++ b/llama_index/storage/index_store/utils.py @@ -1,9 +1,9 @@ from llama_index.constants import DATA_KEY, TYPE_KEY -from llama_index.data_structs.data_structs_v2 import V2IndexStruct +from llama_index.data_structs.data_structs import IndexStruct from llama_index.data_structs.registry import INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS -def index_struct_to_json(index_struct: V2IndexStruct) -> dict: +def index_struct_to_json(index_struct: IndexStruct) -> dict: index_struct_dict = { TYPE_KEY: index_struct.get_type(), DATA_KEY: index_struct.to_dict(), @@ -11,7 +11,7 @@ def index_struct_to_json(index_struct: V2IndexStruct) -> dict: return index_struct_dict -def json_to_index_struct(struct_dict: dict) -> V2IndexStruct: +def json_to_index_struct(struct_dict: dict) -> IndexStruct: type = struct_dict[TYPE_KEY] data_dict = struct_dict[DATA_KEY] cls = INDEX_STRUCT_TYPE_TO_INDEX_STRUCT_CLASS[type] diff --git a/llama_index/vector_stores/chatgpt_plugin.py b/llama_index/vector_stores/chatgpt_plugin.py index e1097e2637e21..99f5edb4668ab 100644 --- a/llama_index/vector_stores/chatgpt_plugin.py +++ b/llama_index/vector_stores/chatgpt_plugin.py @@ -7,7 +7,7 @@ from requests.adapters import HTTPAdapter, Retry from tqdm.auto import tqdm -from llama_index.data_structs.node_v2 import Node, DocumentRelationship +from llama_index.data_structs.node import Node, DocumentRelationship from llama_index.vector_stores.types import ( NodeEmbeddingResult, VectorStore, diff --git a/llama_index/vector_stores/chroma.py b/llama_index/vector_stores/chroma.py index cf13636b91acc..10f446a09d4a3 100644 --- a/llama_index/vector_stores/chroma.py +++ b/llama_index/vector_stores/chroma.py @@ -3,7 +3,7 @@ import math from typing import Any, List, cast -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.utils import truncate_text from llama_index.vector_stores.types import ( NodeEmbeddingResult, diff --git a/llama_index/vector_stores/milvus.py b/llama_index/vector_stores/milvus.py index bdda2a07b3cce..637cd4abd8921 100644 --- a/llama_index/vector_stores/milvus.py +++ b/llama_index/vector_stores/milvus.py @@ -7,7 +7,7 @@ from typing import Any, List, Optional from uuid import uuid4 -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.vector_stores.types import ( NodeEmbeddingResult, VectorStore, diff --git a/llama_index/vector_stores/myscale.py b/llama_index/vector_stores/myscale.py index d2edb636f888e..436ab14f3bd40 100644 --- a/llama_index/vector_stores/myscale.py +++ b/llama_index/vector_stores/myscale.py @@ -7,7 +7,7 @@ import logging from typing import Any, Dict, List, Optional, cast -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.indices.service_context import ServiceContext from llama_index.readers.myscale import ( MyScaleSettings, diff --git a/llama_index/vector_stores/pinecone.py b/llama_index/vector_stores/pinecone.py index b75763ac1741c..5c5d5e233b859 100644 --- a/llama_index/vector_stores/pinecone.py +++ b/llama_index/vector_stores/pinecone.py @@ -10,7 +10,7 @@ from functools import partial from typing import Any, Callable, Dict, List, Optional, cast -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.vector_stores.types import ( NodeEmbeddingResult, VectorStore, diff --git a/llama_index/vector_stores/qdrant.py b/llama_index/vector_stores/qdrant.py index 2a08305b046e7..7f767057ed1c7 100644 --- a/llama_index/vector_stores/qdrant.py +++ b/llama_index/vector_stores/qdrant.py @@ -6,7 +6,7 @@ import logging from typing import Any, List, Optional, cast -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.utils import iter_batch from llama_index.vector_stores.types import ( NodeEmbeddingResult, diff --git a/llama_index/vector_stores/types.py b/llama_index/vector_stores/types.py index 87bb6e874dc51..7891f67f7ac39 100644 --- a/llama_index/vector_stores/types.py +++ b/llama_index/vector_stores/types.py @@ -5,7 +5,7 @@ from typing import Any, List, Optional, Protocol, runtime_checkable from enum import Enum -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node DEFAULT_PERSIST_DIR = "./storage" diff --git a/tests/indices/conftest.py b/tests/indices/conftest.py index 3ff9225025f9b..c1d29798f8c78 100644 --- a/tests/indices/conftest.py +++ b/tests/indices/conftest.py @@ -1,6 +1,6 @@ from typing import List import pytest -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.readers.schema.base import Document diff --git a/tests/indices/empty/test_base.py b/tests/indices/empty/test_base.py index d2170df7bbe9e..7a9d07af5b868 100644 --- a/tests/indices/empty/test_base.py +++ b/tests/indices/empty/test_base.py @@ -1,6 +1,6 @@ """Test empty index.""" -from llama_index.data_structs.data_structs_v2 import EmptyIndex +from llama_index.data_structs.data_structs import EmptyIndex from llama_index.indices.empty.base import GPTEmptyIndex from llama_index.indices.service_context import ServiceContext diff --git a/tests/indices/knowledge_graph/test_base.py b/tests/indices/knowledge_graph/test_base.py index eff7a2d018c29..8107a7abf94bb 100644 --- a/tests/indices/knowledge_graph/test_base.py +++ b/tests/indices/knowledge_graph/test_base.py @@ -4,7 +4,7 @@ from unittest.mock import patch import pytest -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.embeddings.base import BaseEmbedding from llama_index.indices.knowledge_graph.base import GPTKnowledgeGraphIndex from llama_index.indices.service_context import ServiceContext diff --git a/tests/indices/list/test_index.py b/tests/indices/list/test_index.py index 4c75711fd0a49..38c9ba52a2977 100644 --- a/tests/indices/list/test_index.py +++ b/tests/indices/list/test_index.py @@ -2,7 +2,7 @@ from typing import Dict, List, Tuple -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.indices.base_retriever import BaseRetriever from llama_index.indices.list.base import GPTListIndex, ListRetrieverMode from llama_index.indices.service_context import ServiceContext diff --git a/tests/indices/postprocessor/test_base.py b/tests/indices/postprocessor/test_base.py index fe2c88f9f9ba3..cd15487a294a0 100644 --- a/tests/indices/postprocessor/test_base.py +++ b/tests/indices/postprocessor/test_base.py @@ -7,7 +7,7 @@ from llama_index.indices.query.schema import QueryBundle from llama_index.prompts.prompts import Prompt, SimpleInputPrompt from llama_index.indices.service_context import ServiceContext -from llama_index.data_structs.node_v2 import Node, DocumentRelationship, NodeWithScore +from llama_index.data_structs.node import Node, DocumentRelationship, NodeWithScore from llama_index.indices.postprocessor.node import ( PrevNextNodePostprocessor, KeywordNodePostprocessor, diff --git a/tests/indices/query/test_compose_vector.py b/tests/indices/query/test_compose_vector.py index 8966e2614e648..3e5d2cff91a7c 100644 --- a/tests/indices/query/test_compose_vector.py +++ b/tests/indices/query/test_compose_vector.py @@ -7,7 +7,7 @@ import pytest -from llama_index.data_structs.data_structs_v2 import V2IndexStruct +from llama_index.data_structs.data_structs import IndexStruct from llama_index.embeddings.base import BaseEmbedding from llama_index.indices.composability.graph import ComposableGraph from llama_index.indices.keyword_table.simple_base import GPTSimpleKeywordTableIndex @@ -146,8 +146,8 @@ def test_recursive_query_vector_table_query_configs( vector2 = GPTVectorStoreIndex.from_documents( documents[2:4], service_context=mock_service_context, **vector_kwargs ) - assert isinstance(vector1.index_struct, V2IndexStruct) - assert isinstance(vector2.index_struct, V2IndexStruct) + assert isinstance(vector1.index_struct, IndexStruct) + assert isinstance(vector2.index_struct, IndexStruct) vector1.index_struct.index_id = "vector1" vector2.index_struct.index_id = "vector2" summaries = [ diff --git a/tests/indices/struct_store/test_base.py b/tests/indices/struct_store/test_base.py index 48be9ac3d38f2..934f626b2dd86 100644 --- a/tests/indices/struct_store/test_base.py +++ b/tests/indices/struct_store/test_base.py @@ -28,7 +28,7 @@ from tests.mock_utils.mock_prompts import ( MOCK_TABLE_CONTEXT_PROMPT, ) -from llama_index.data_structs.node_v2 import Node, DocumentRelationship +from llama_index.data_structs.node import Node, DocumentRelationship def _delete_table_items(engine: Any, table: Table) -> None: diff --git a/tests/indices/test_loading.py b/tests/indices/test_loading.py index be0614d90aeea..6fe955cef0326 100644 --- a/tests/indices/test_loading.py +++ b/tests/indices/test_loading.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import List import pytest -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.indices.list.base import GPTListIndex from llama_index.indices.loading import ( load_index_from_storage, diff --git a/tests/indices/test_prompt_helper.py b/tests/indices/test_prompt_helper.py index f0ec503c33036..e7f31ce50dc92 100644 --- a/tests/indices/test_prompt_helper.py +++ b/tests/indices/test_prompt_helper.py @@ -3,7 +3,7 @@ from langchain import PromptTemplate as LangchainPrompt -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.indices.prompt_helper import PromptHelper from llama_index.prompts.base import Prompt from tests.mock_utils.mock_utils import mock_tokenizer diff --git a/tests/indices/tree/test_embedding_retriever.py b/tests/indices/tree/test_embedding_retriever.py index a4e79407fdf51..de96dd67981d0 100644 --- a/tests/indices/tree/test_embedding_retriever.py +++ b/tests/indices/tree/test_embedding_retriever.py @@ -6,7 +6,7 @@ import pytest -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.indices.query.schema import QueryBundle from llama_index.indices.service_context import ServiceContext from llama_index.indices.tree.select_leaf_embedding_retriever import ( diff --git a/tests/indices/tree/test_index.py b/tests/indices/tree/test_index.py index aadf0f40902c2..cf6a8726e96cc 100644 --- a/tests/indices/tree/test_index.py +++ b/tests/indices/tree/test_index.py @@ -3,8 +3,8 @@ from typing import Any, Dict, List, Optional from unittest.mock import patch -from llama_index.data_structs.data_structs_v2 import IndexGraph -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.data_structs import IndexGraph +from llama_index.data_structs.node import Node from llama_index.indices.service_context import ServiceContext from llama_index.storage.docstore import BaseDocumentStore from llama_index.indices.tree.base import GPTTreeIndex diff --git a/tests/indices/vector_store/test_faiss.py b/tests/indices/vector_store/test_faiss.py index 4057c3c8e479d..7b2bb6a67f4fe 100644 --- a/tests/indices/vector_store/test_faiss.py +++ b/tests/indices/vector_store/test_faiss.py @@ -5,7 +5,7 @@ from typing import List import pytest -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.indices.service_context import ServiceContext from llama_index.indices.vector_store.base import GPTVectorStoreIndex diff --git a/tests/indices/vector_store/test_myscale.py b/tests/indices/vector_store/test_myscale.py index e121fbd5e7f88..6ebc44ea465ec 100644 --- a/tests/indices/vector_store/test_myscale.py +++ b/tests/indices/vector_store/test_myscale.py @@ -12,7 +12,7 @@ except ImportError: clickhouse_connect = None # type: ignore -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.readers.schema.base import Document from llama_index.vector_stores import MyScaleVectorStore from llama_index.vector_stores.types import VectorStoreQuery diff --git a/tests/indices/vector_store/test_retrievers.py b/tests/indices/vector_store/test_retrievers.py index 45bb3690379be..a5b132e9af85e 100644 --- a/tests/indices/vector_store/test_retrievers.py +++ b/tests/indices/vector_store/test_retrievers.py @@ -1,5 +1,5 @@ from typing import List, cast -from llama_index.data_structs.node_v2 import DocumentRelationship, Node +from llama_index.data_structs.node import DocumentRelationship, Node from llama_index.indices.query.schema import QueryBundle from llama_index.indices.service_context import ServiceContext from llama_index.indices.vector_store.base import GPTVectorStoreIndex diff --git a/tests/storage/docstore/test_simple_docstore.py b/tests/storage/docstore/test_simple_docstore.py index 5e5cf568ec585..241cfce1ad971 100644 --- a/tests/storage/docstore/test_simple_docstore.py +++ b/tests/storage/docstore/test_simple_docstore.py @@ -3,7 +3,7 @@ from pathlib import Path import pytest -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.storage.docstore import SimpleDocumentStore from llama_index.readers.schema.base import Document from llama_index.storage.kvstore.simple_kvstore import SimpleKVStore diff --git a/tests/vector_stores/test_weaviate.py b/tests/vector_stores/test_weaviate.py index fe0d86986555a..05d1b4ebc778b 100644 --- a/tests/vector_stores/test_weaviate.py +++ b/tests/vector_stores/test_weaviate.py @@ -1,6 +1,6 @@ import sys from unittest.mock import MagicMock -from llama_index.data_structs.node_v2 import Node +from llama_index.data_structs.node import Node from llama_index.vector_stores.types import NodeEmbeddingResult from llama_index.vector_stores.weaviate import WeaviateVectorStore