"""Base index classes.""" import json import logging from abc import abstractmethod from typing import ( Any, Dict, Generic, List, Optional, Sequence, Type, TypeVar, Union, cast, ) from gpt_index.data_structs.data_structs import IndexStruct, Node from gpt_index.docstore import DOC_TYPE, DocumentStore from gpt_index.embeddings.base import BaseEmbedding from gpt_index.embeddings.openai import OpenAIEmbedding from gpt_index.indices.node_utils import get_nodes_from_document from gpt_index.indices.prompt_helper import PromptHelper from gpt_index.indices.query.base import BaseGPTIndexQuery from gpt_index.indices.query.query_runner import QueryRunner from gpt_index.indices.query.query_transform import BaseQueryTransform from gpt_index.indices.query.schema import QueryBundle, QueryConfig, QueryMode from gpt_index.indices.registry import IndexRegistry from gpt_index.langchain_helpers.chain_wrapper import LLMPredictor from gpt_index.langchain_helpers.text_splitter import TextSplitter, TokenTextSplitter from gpt_index.readers.schema.base import Document from gpt_index.response.schema import Response from gpt_index.schema import BaseDocument from gpt_index.token_counter.token_counter import llm_token_counter IS = TypeVar("IS", bound=IndexStruct) DOCUMENTS_INPUT = Union[BaseDocument, "BaseGPTIndex"] class BaseGPTIndex(Generic[IS]): """Base LlamaIndex. Args: documents (Optional[Sequence[BaseDocument]]): List of documents to build the index from. llm_predictor (LLMPredictor): Optional LLMPredictor object. If not provided, will use the default LLMPredictor (text-davinci-003) prompt_helper (PromptHelper): Optional PromptHelper object. If not provided, will use the default PromptHelper. chunk_size_limit (Optional[int]): Optional chunk size limit. If not provided, will use the default chunk size limit (4096 max input size). include_extra_info (bool): Optional bool. If True, extra info (i.e. metadata) of each Document will be prepended to its text to help with queries. Default is True. """ index_struct_cls: Type[IS] def __init__( self, documents: Optional[Sequence[DOCUMENTS_INPUT]] = None, index_struct: Optional[IS] = None, llm_predictor: Optional[LLMPredictor] = None, embed_model: Optional[BaseEmbedding] = None, docstore: Optional[DocumentStore] = None, index_registry: Optional[IndexRegistry] = None, prompt_helper: Optional[PromptHelper] = None, text_splitter: Optional[TextSplitter] = None, chunk_size_limit: Optional[int] = None, include_extra_info: bool = True, ) -> None: """Initialize with parameters.""" if index_struct is None and documents is None: raise ValueError("One of documents or index_struct must be provided.") if index_struct is not None and documents is not None: raise ValueError("Only one of documents or index_struct can be provided.") self._llm_predictor = llm_predictor or LLMPredictor() # NOTE: the embed_model isn't used in all indices self._embed_model = embed_model or OpenAIEmbedding() self._include_extra_info = include_extra_info # TODO: move out of base if we need custom params per index self._prompt_helper = prompt_helper or PromptHelper.from_llm_predictor( self._llm_predictor, chunk_size_limit=chunk_size_limit ) self._text_splitter = text_splitter or self._build_fallback_text_splitter() # build index struct in the init function self._docstore = docstore or DocumentStore() 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) documents = self._process_documents( documents, self._docstore, self._index_registry ) self._validate_documents(documents) # TODO: introduce document store outside __init__ function self._index_struct = self.build_index_from_documents(documents) # update index registry and docstore with index_struct self._update_index_registry_and_docstore() @property def prompt_helper(self) -> PromptHelper: """Get the prompt helper corresponding to the index.""" return self._prompt_helper @property def docstore(self) -> DocumentStore: """Get the docstore corresponding to the index.""" return self._docstore @property def index_registry(self) -> IndexRegistry: """Get the index registry corresponding to the index.""" return self._index_registry @property def llm_predictor(self) -> LLMPredictor: """Get the llm predictor.""" return self._llm_predictor @property def embed_model(self) -> BaseEmbedding: """Get the llm predictor.""" return self._embed_model def _update_index_registry_and_docstore(self) -> None: """Update index registry and docstore.""" # update index registry with current struct cur_type = self.index_struct_cls.get_type() self._index_registry.type_to_struct[cur_type] = self.index_struct_cls self._index_registry.type_to_query[cur_type] = self.get_query_map() # update docstore with current struct # NOTE: we call allow_update=True: in old versions of the docstore, # the index_struct was not stored in the docstore. whereas # in the new docstore, index_struct is stored in the docstore. # if we want to break BW compatibility, we can just remove this line # and only insert into docstore during index construction. self._docstore.add_documents([self.index_struct], allow_update=True) def _process_documents( self, documents: Sequence[DOCUMENTS_INPUT], docstore: DocumentStore, index_registry: IndexRegistry, ) -> List[BaseDocument]: """Process documents.""" results: List[DOC_TYPE] = [] for doc in documents: if isinstance(doc, BaseGPTIndex): # if user passed in another index, we need to do the following: # - update docstore with the docstore in the index # - validate that the index is in the docstore # - update the index registry index_registry.update(doc.index_registry) docstore.update_docstore(doc.docstore) # assert that the doc exists within the docstore sub_index_struct = doc.index_struct_with_text if not docstore.document_exists(sub_index_struct.get_doc_id()): raise ValueError( "The index struct of the sub-index must exist in the docstore. " f"Invalid doc ID: {sub_index_struct.get_doc_id()}" ) results.append(sub_index_struct) elif isinstance(doc, Document): results.append(doc) else: raise ValueError(f"Invalid document type: {type(doc)}.") return cast(List[BaseDocument], results) def _validate_documents(self, documents: Sequence[BaseDocument]) -> None: """Validate documents.""" for doc in documents: if not isinstance(doc, BaseDocument): raise ValueError("Documents must be of type BaseDocument.") @property def index_struct(self) -> IS: """Get the index struct.""" return self._index_struct @property def index_struct_with_text(self) -> IS: """Get the index struct with text. If text not set, raise an error. For use when composing indices with other indices. """ # make sure that we generate text for index struct if self._index_struct.text is None: # NOTE: set text to be empty string for now raise ValueError( "Index must have text property set in order " "to be composed with other indices. " "In order to set text, please run `index.set_text()`." ) return self._index_struct def set_text(self, text: str) -> None: """Set summary text for index struct. This allows index_struct_with_text to be used to compose indices with other indices. """ self._index_struct.text = text def set_extra_info(self, extra_info: Dict[str, Any]) -> None: """Set extra info (metadata) for index struct. If this index is used as a subindex for a parent index, the metadata will be propagated to all nodes derived from this subindex, in the parent index. """ self._index_struct.extra_info = extra_info def set_doc_id(self, doc_id: str) -> None: """Set doc_id for index struct. This is used to uniquely identify the index struct in the docstore. 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], allow_update=True) def get_doc_id(self) -> str: """Get doc_id for index struct. If doc_id not set, raise an error. """ if self._index_struct.doc_id is None: raise ValueError("Index must have doc_id property set.") return self._index_struct.doc_id def _get_nodes_from_document( self, document: BaseDocument, start_idx: int = 0, ) -> List[Node]: return get_nodes_from_document( document=document, text_splitter=self._text_splitter, start_idx=start_idx, include_extra_info=self._include_extra_info, ) def _build_fallback_text_splitter(self) -> TextSplitter: """Build the text splitter if not specified in args.""" return TokenTextSplitter() @abstractmethod def _build_index_from_documents(self, documents: Sequence[BaseDocument]) -> IS: """Build the index from documents.""" @llm_token_counter("build_index_from_documents") def build_index_from_documents(self, documents: Sequence[BaseDocument]) -> IS: """Build the index from documents.""" return self._build_index_from_documents(documents) @abstractmethod def _insert(self, document: BaseDocument, **insert_kwargs: Any) -> None: """Insert a document.""" @llm_token_counter("insert") def insert(self, document: DOCUMENTS_INPUT, **insert_kwargs: Any) -> None: """Insert a document. Args: document (Union[BaseDocument, BaseGPTIndex]): document to insert """ processed_doc = self._process_documents( [document], self._docstore, self._index_registry )[0] self._validate_documents([processed_doc]) self._insert(processed_doc, **insert_kwargs) @abstractmethod def _delete(self, doc_id: str, **delete_kwargs: Any) -> None: """Delete a document.""" def delete(self, doc_id: str, **delete_kwargs: Any) -> None: """Delete a document from the index. All nodes in the index related to the index will be deleted. Args: doc_id (str): document id """ logging.debug(f"> Deleting document: {doc_id}") self._delete(doc_id, **delete_kwargs) def update(self, document: DOCUMENTS_INPUT, **update_kwargs: Any) -> None: """Update a document. This is equivalent to deleting the document and then inserting it again. Args: document (Union[BaseDocument, BaseGPTIndex]): document to update insert_kwargs (Dict): kwargs to pass to insert delete_kwargs (Dict): kwargs to pass to delete """ self.delete(document.get_doc_id(), **update_kwargs.pop("delete_kwargs", {})) self.insert(document, **update_kwargs.pop("insert_kwargs", {})) def _preprocess_query(self, mode: QueryMode, query_kwargs: Dict) -> None: """Preprocess query. This allows subclasses to pass in additional query kwargs to query, for instance arguments that are shared between the index and the query class. By default, this does nothing. This also allows subclasses to do validation. """ pass def query( self, query_str: Union[str, QueryBundle], mode: str = QueryMode.DEFAULT, query_transform: Optional[BaseQueryTransform] = None, use_async: bool = False, **query_kwargs: Any, ) -> Response: """Answer a query. When `query` is called, we query the index with the given `mode` and `query_kwargs`. The `mode` determines the type of query to run, and `query_kwargs` are parameters that are specific to the query type. For a comprehensive documentation of available `mode` and `query_kwargs` to query a given index, please visit :ref:`Ref-Query`. """ mode_enum = QueryMode(mode) if mode_enum == QueryMode.RECURSIVE: # TODO: deprecated, use ComposableGraph instead. if "query_configs" not in query_kwargs: raise ValueError("query_configs must be provided for recursive mode.") query_configs = query_kwargs["query_configs"] query_runner = QueryRunner( self._llm_predictor, self._prompt_helper, self._embed_model, self._docstore, self._index_registry, query_configs=query_configs, query_transform=query_transform, recursive=True, use_async=use_async, ) return query_runner.query(query_str, self._index_struct) else: self._preprocess_query(mode_enum, query_kwargs) # TODO: pass in query config directly query_config = QueryConfig( index_struct_type=self._index_struct.get_type(), query_mode=mode_enum, query_kwargs=query_kwargs, ) query_runner = QueryRunner( self._llm_predictor, self._prompt_helper, self._embed_model, self._docstore, self._index_registry, query_configs=[query_config], query_transform=query_transform, recursive=False, use_async=use_async, ) return query_runner.query(query_str, self._index_struct) async def aquery( self, query_str: Union[str, QueryBundle], mode: str = QueryMode.DEFAULT, query_transform: Optional[BaseQueryTransform] = None, **query_kwargs: Any, ) -> Response: """Asynchronously answer a query. When `query` is called, we query the index with the given `mode` and `query_kwargs`. The `mode` determines the type of query to run, and `query_kwargs` are parameters that are specific to the query type. For a comprehensive documentation of available `mode` and `query_kwargs` to query a given index, please visit :ref:`Ref-Query`. """ # TODO: currently we don't have async versions of all # underlying functions. Setting use_async=True # will cause async nesting errors because we assume # it's called in a synchronous setting. use_async = False mode_enum = QueryMode(mode) if mode_enum == QueryMode.RECURSIVE: # TODO: deprecated, use ComposableGraph instead. if "query_configs" not in query_kwargs: raise ValueError("query_configs must be provided for recursive mode.") query_configs = query_kwargs["query_configs"] query_runner = QueryRunner( self._llm_predictor, self._prompt_helper, self._embed_model, self._docstore, self._index_registry, query_configs=query_configs, query_transform=query_transform, recursive=True, use_async=use_async, ) return await query_runner.aquery(query_str, self._index_struct) else: self._preprocess_query(mode_enum, query_kwargs) # TODO: pass in query config directly query_config = QueryConfig( index_struct_type=self._index_struct.get_type(), query_mode=mode_enum, query_kwargs=query_kwargs, ) query_runner = QueryRunner( self._llm_predictor, self._prompt_helper, self._embed_model, self._docstore, self._index_registry, query_configs=[query_config], query_transform=query_transform, recursive=False, use_async=use_async, ) return await query_runner.aquery(query_str, self._index_struct) @classmethod @abstractmethod def get_query_map(cls) -> Dict[str, Type[BaseGPTIndexQuery]]: """Get query map.""" @classmethod def load_from_dict( cls, result_dict: Dict[str, Any], **kwargs: Any ) -> "BaseGPTIndex": """Load index from dict.""" if "index_struct" in result_dict: index_struct = cls.index_struct_cls.from_dict(result_dict["index_struct"]) index_struct_id = index_struct.get_doc_id() elif "index_struct_id" in result_dict: index_struct_id = result_dict["index_struct_id"] else: raise ValueError("index_struct or index_struct_id must be provided.") type_to_struct = {cls.index_struct_cls.get_type(): cls.index_struct_cls} # NOTE: index_struct can have multiple types for backwards compatibility, # map to same class type_to_struct = { index_type: cls.index_struct_cls for index_type in cls.index_struct_cls.get_types() } docstore = DocumentStore.load_from_dict( result_dict["docstore"], type_to_struct=type_to_struct, ) if "index_struct_id" in result_dict: index_struct = docstore.get_document(index_struct_id) return cls(index_struct=index_struct, docstore=docstore, **kwargs) @classmethod def load_from_string(cls, index_string: str, **kwargs: Any) -> "BaseGPTIndex": """Load index from string (in JSON-format). This method loads the index from a JSON string. The index data structure itself is preserved completely. If the index is defined over subindices, those subindices will also be preserved (and subindices of those subindices, etc.). NOTE: load_from_string should not be used for indices composed on top of other indices. Please define a `ComposableGraph` and use `save_to_string` and `load_from_string` on that instead. Args: index_string (str): The index string (in JSON-format). Returns: BaseGPTIndex: The loaded index. """ result_dict = json.loads(index_string) return cls.load_from_dict(result_dict, **kwargs) @classmethod def load_from_disk(cls, save_path: str, **kwargs: Any) -> "BaseGPTIndex": """Load index from disk. This method loads the index from a JSON file stored on disk. The index data structure itself is preserved completely. If the index is defined over subindices, those subindices will also be preserved (and subindices of those subindices, etc.). NOTE: load_from_disk should not be used for indices composed on top of other indices. Please define a `ComposableGraph` and use `save_to_disk` and `load_from_disk` on that instead. Args: save_path (str): The save_path of the file. Returns: BaseGPTIndex: The loaded index. """ with open(save_path, "r") as f: file_contents = f.read() return cls.load_from_string(file_contents, **kwargs) def save_to_dict(self, **save_kwargs: Any) -> dict: """Save to dict.""" if self.docstore.contains_index_struct( exclude_ids=[self.index_struct.get_doc_id()] ): raise ValueError( "Cannot call save index if index is composed on top of " "other indices. Please define a `ComposableGraph` and use " "`save_to_string` and `load_from_string` on that instead." ) out_dict: Dict[str, Any] = { "index_struct_id": self.index_struct.get_doc_id(), "docstore": self.docstore.serialize_to_dict(), } return out_dict def save_to_string(self, **save_kwargs: Any) -> str: """Save to string. This method stores the index into a JSON string. NOTE: save_to_string should not be used for indices composed on top of other indices. Please define a `ComposableGraph` and use `save_to_string` and `load_from_string` on that instead. Returns: str: The JSON string of the index. """ out_dict = self.save_to_dict(**save_kwargs) return json.dumps(out_dict, **save_kwargs) def save_to_disk(self, save_path: str, **save_kwargs: Any) -> None: """Save to file. This method stores the index into a JSON file stored on disk. NOTE: save_to_disk should not be used for indices composed on top of other indices. Please define a `ComposableGraph` and use `save_to_disk` and `load_from_disk` on that instead. Args: save_path (str): The save_path of the file. """ index_string = self.save_to_string(**save_kwargs) with open(save_path, "w") as f: f.write(index_string)