MongoDBQueryEngine

def connect_db(self, *args: Any, **kwargs: Any) -> bool:
    """
    Connects to the MongoDB database and initializes the query index from the existing collection.

    This method verifies the existence of the collection, sets up the database connection,
    builds the vector store index, and pings the MongoDB server.

    Returns:
        bool: True if connection is successful; False otherwise.
    """
    try:
        # Check if the target collection exists.
        if not self._check_existing_collection():
            raise ValueError(
                f"Collection '{self.collection_name}' not found in database '{self.database_name}'. "
                "Please run init_db to create a new collection."
            )
        # Reinitialize without overwriting the existing collection.
        self._set_up(overwrite=False)

        self.index = VectorStoreIndex.from_vector_store(
            vector_store=self.vector_search_engine,  # type: ignore[arg-type]
            storage_context=self.storage_context,
            embed_model=self.embedding_model,
        )

        self.vector_db.client.admin.command("ping")  # type: ignore[union-attr]
        logger.info("Connected to MongoDB successfully.")
        return True
    except Exception as error:
        logger.error("Failed to connect to MongoDB: %s", error)
        return False

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:
    """
    Initializes the MongoDB database by creating or overwriting the collection and indexing documents.

    This method loads documents from a directory or provided file paths, sets up the database (optionally
    overwriting any existing collection), builds the vector store index, and inserts the documents.

    Args:
        new_doc_dir (Optional[Union[Path, str]]): Directory containing documents to be indexed.
        new_doc_paths_or_urls (Optional[Sequence[Union[Path, str]]]): List of file paths or URLs for documents.
        *args (Any): Additional positional arguments.
        **kwargs (Any): Additional keyword arguments.

    Returns:
        bool: True if the database is successfully initialized; False otherwise.
    """
    try:
        # Check if the collection already exists.
        if self._check_existing_collection():
            logger.warning(
                f"Collection '{self.collection_name}' already exists in database '{self.database_name}'. "
                "Please use connect_db to connect to the existing collection or use init_db to overwrite it."
            )
        # Set up the database with overwriting.
        self._set_up(overwrite=True)
        self.vector_db.client.admin.command("ping")  # type: ignore[union-attr]
        # Gather document paths.
        logger.info("Setting up the database with existing collection.")
        documents = self._load_doc(input_dir=new_doc_dir, input_docs=new_doc_paths_or_urls)
        self.index = VectorStoreIndex.from_vector_store(
            vector_store=self.vector_search_engine,  # type: ignore[arg-type]
            storage_context=self.storage_context,
            embed_model=self.embedding_model,
        )
        for doc in documents:
            self.index.insert(doc)
        logger.info("Database initialized with %d documents.", len(documents))
        return True
    except Exception as e:
        logger.error("Failed to initialize the database: %s", e)
        return False

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:
    """
    Adds new documents to the existing vector store index.

    This method validates that the index exists, loads documents from the specified directory or file paths,
    and inserts them into the vector store index.

    Args:
        new_doc_dir (Optional[Union[Path, str]]): Directory containing new documents.
        new_doc_paths_or_urls (Optional[Sequence[Union[Path, str]]]): List of file paths or URLs for new documents.
        *args (Any): Additional positional 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)  # type: ignore[union-attr]

def query(self, question: str, *args: Any, **kwargs: Any) -> Any:  # type: ignore[no-any-unimported, type-arg]
    """
    Queries the indexed documents using the provided question.

    This method validates that the query index is initialized, creates a query engine from the vector store index,
    and executes the query. If the response is empty, a default reply is returned.

    Args:
        question (str): The query question.
        args (Any): Additional positional arguments.
        kwargs (Any): Additional keyword arguments.

    Returns:
        Any: The query response as a string, or a default reply if no results are found.
    """
    self._validate_query_index()
    self.query_engine = self.index.as_query_engine(llm=self.llm)  # type: ignore[union-attr]
    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:
    """
    Retrieves the name of the MongoDB collection.

    Returns:
        str: The collection name.

    Raises:
        ValueError: If the collection name is not set.
    """
    if self.collection_name:
        return self.collection_name
    else:
        raise ValueError("Collection name not set.")