[Question]: Can i use multiple collections in mongo at a time?

[arnavroh45]

Question Validation

• I have searched both the documentation and discord for an answer.

Question

I want to use multiple collections from Mongo to retrieve data. Is it possible to do so?
This is my current approach.

store = MongoDBAtlasVectorSearch(client, db_name, collection_name, index_name)
index = VectorStoreIndex.from_vector_store(vector_store=store)
query_engine = index.as_query_engine()

[dosubot]

Yes, you can use multiple collections from MongoDB to retrieve data. To achieve this, you can extend the MongoDBAtlasVectorSearch class to support multiple collections. Here is an example of how you might modify the class:

class MongoDBAtlasVectorSearch(VectorStore):
    def __init__(
        self,
        mongodb_client: Optional[Any] = None,
        db_name: str = "default_db",
        collection_names: List[str] = ["default_collection"],
        index_name: str = "default",
        id_key: str = "id",
        embedding_key: str = "embedding",
        text_key: str = "text",
        metadata_key: str = "metadata",
        insert_kwargs: Optional[Dict] = None,
        **kwargs: Any,
    ) -> None:
        import_err_msg = "`pymongo` package not found, please run `pip install pymongo`"
        try:
            from importlib.metadata import version
            from pymongo import MongoClient
            from pymongo.driver_info import DriverInfo
        except ImportError:
            raise ImportError(import_err_msg)

        if mongodb_client is not None:
            self._mongodb_client = cast(MongoClient, mongodb_client)
        else:
            if "MONGO_URI" not in os.environ:
                raise ValueError(
                    "Must specify MONGO_URI via env variable "
                    "if not directly passing in client."
                )
            self._mongodb_client = MongoClient(
                os.environ["MONGO_URI"],
                driver=DriverInfo(name="llama-index", version=version("llama-index")),
            )

        self._collections = {name: self._mongodb_client[db_name][name] for name in collection_names}
        self._current_collection = self._collections[collection_names[0]]
        self._index_name = index_name
        self._embedding_key = embedding_key
        self._id_key = id_key
        self._text_key = text_key
        self._metadata_key = metadata_key
        self._insert_kwargs = insert_kwargs or {}

    def switch_collection(self, collection_name: str) -> None:
        if collection_name in self._collections:
            self._current_collection = self._collections[collection_name]
        else:
            raise ValueError(f"Collection {collection_name} not found.")

    # Modify methods to use self._current_collection instead of self._collection

This modification allows you to switch between multiple collections using the switch_collection method. You would also need to adjust your methods to use self._current_collection instead of self._collection ^[1][2].

To continue talking to Dosu, mention @dosu.