ChromaDBQueryEngine

def init_db(
    self,
    new_doc_dir: Optional[Union[Path, str]] = None,
    new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
    *args: Any,
    **kwargs: Any,
) -> bool:
    """Initialize the database with the input documents or records.
    It overwrites the existing collection in the database.

    It takes the following steps,
    1. Set up ChromaDB and LlamaIndex storage.
    2. insert documents and build indexes upon them.

    Args:
        new_doc_dir: a dir of input documents that are used to create the records in database.
        new_doc_paths_or_urls:
            a sequence of input documents that are used to create the records in database.
            a document can be a path to a file or a url.
        *args: Any additional arguments
        **kwargs: Any additional keyword arguments

    Returns:
        bool: True if initialization is successful

    """

    self._set_up(overwrite=True)
    documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
    self.index = VectorStoreIndex.from_documents(documents=documents, storage_context=self.storage_context)
    return True

def connect_db(self, *args: Any, **kwargs: Any) -> bool:
    """Connect to the database.
    It does not overwrite the existing collection in the database.

     It takes the following steps,
    1. Set up ChromaDB and LlamaIndex storage.
    2. Create the llamaIndex vector store index for querying or inserting docs later

    Args:
        *args: Any additional arguments
        **kwargs: Any additional keyword arguments

    Returns:
        bool: True if connection is successful
    """
    self._set_up(overwrite=False)
    self.index = VectorStoreIndex.from_vector_store(
        vector_store=self.vector_store, storage_context=self.storage_context
    )

    return True

def add_docs(
    self,
    new_doc_dir: Optional[Union[Path, str]] = None,
    new_doc_paths_or_urls: Optional[Sequence[Union[Path, str]]] = None,
    *args: Any,
    **kwargs: Any,
) -> None:
    """Add new documents to the underlying database and add to the index.

    Args:
        new_doc_dir: A dir of input documents that are used to create the records in database.
        new_doc_paths_or_urls: A sequence of input documents that are used to create the records in database. A document can be a path to a file or a url.
        *args: Any additional arguments
        **kwargs: Any additional keyword arguments
    """
    self._validate_query_index()
    documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
    for doc in documents:
        self.index.insert(doc)

def query(self, question: str) -> str:
    """
    Retrieve information from indexed documents by processing a query using the engine's LLM.

    Args:
        question: A natural language query string used to search the indexed documents.

    Returns:
        A string containing the response generated by LLM.
    """
    self._validate_query_index()
    self.query_engine = self.index.as_query_engine(llm=self.llm)
    response = self.query_engine.query(question)

    if str(response) == EMPTY_RESPONSE_TEXT:
        return EMPTY_RESPONSE_REPLY

    return str(response)

def get_collection_name(self) -> str:
    """
    Get the name of the collection used by the query engine.

    Returns:
        The name of the collection.
    """
    if self.collection_name:
        return self.collection_name
    else:
        raise ValueError("Collection name not set.")