keyval_docstore and mongodb docstore updated to accept custom collections suffix (#11873)

llama-index-core/llama_index/core/storage/docstore/keyval_docstore.py

@@ -11,6 +11,9 @@ from llama_index.core.storage.docstore.utils import doc_to_json, json_to_doc
 from llama_index.core.storage.kvstore.types import DEFAULT_BATCH_SIZE, BaseKVStore
 
 DEFAULT_NAMESPACE = "docstore"
+DEFAULT_COLLECTION_DATA_SUFFIX = "/data"
+DEFAULT_REF_DOC_COLLECTION_SUFFIX = "/ref_doc_info"
+DEFAULT_METADATA_COLLECTION_SUFFIX = "/metadata"
 
 
 class KVDocumentStore(BaseDocumentStore):
@@ -46,13 +49,27 @@ class KVDocumentStore(BaseDocumentStore):
         kvstore: BaseKVStore,
         namespace: Optional[str] = None,
         batch_size: int = DEFAULT_BATCH_SIZE,
+        node_collection_suffix: Optional[str] = None,
+        ref_doc_collection_suffix: Optional[str] = None,
+        metadata_collection_suffix: Optional[str] = None,
     ) -> None:
         """Init a KVDocumentStore."""
         self._kvstore = kvstore
         self._namespace = namespace or DEFAULT_NAMESPACE
-        self._node_collection = f"{self._namespace}/data"
-        self._ref_doc_collection = f"{self._namespace}/ref_doc_info"
-        self._metadata_collection = f"{self._namespace}/metadata"
+        self._node_collection_suffix = (
+            node_collection_suffix or DEFAULT_COLLECTION_DATA_SUFFIX
+        )
+        self._ref_doc_collection_suffix = (
+            ref_doc_collection_suffix or DEFAULT_REF_DOC_COLLECTION_SUFFIX
+        )
+        self._metadata_collection_suffix = (
+            metadata_collection_suffix or DEFAULT_METADATA_COLLECTION_SUFFIX
+        )
+        self._node_collection = f"{self._namespace}{self._node_collection_suffix}"
+        self._ref_doc_collection = f"{self._namespace}{self._ref_doc_collection_suffix}"
+        self._metadata_collection = (
+            f"{self._namespace}{self._metadata_collection_suffix}"
+        )
         self._batch_size = batch_size
 
     @property

llama-index-integrations/storage/docstore/llama-index-storage-docstore-mongodb/llama_index/storage/docstore/mongodb/base.py

@@ -20,10 +20,20 @@ class MongoDocumentStore(KVDocumentStore):
         self,
         mongo_kvstore: MongoDBKVStore,
         namespace: Optional[str] = None,
+        node_collection_suffix: Optional[str] = None,
+        ref_doc_collection_suffix: Optional[str] = None,
+        metadata_collection_suffix: Optional[str] = None,
         batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> None:
         """Init a MongoDocumentStore."""
-        super().__init__(mongo_kvstore, namespace=namespace, batch_size=batch_size)
+        super().__init__(
+            mongo_kvstore,
+            namespace=namespace,
+            batch_size=batch_size,
+            node_collection_suffix=node_collection_suffix,
+            ref_doc_collection_suffix=ref_doc_collection_suffix,
+            metadata_collection_suffix=metadata_collection_suffix,
+        )
 
     @classmethod
     def from_uri(
@@ -31,10 +41,19 @@ class MongoDocumentStore(KVDocumentStore):
         uri: str,
         db_name: Optional[str] = None,
         namespace: Optional[str] = None,
+        node_collection_suffix: Optional[str] = None,
+        ref_doc_collection_suffix: Optional[str] = None,
+        metadata_collection_suffix: Optional[str] = None,
     ) -> "MongoDocumentStore":
         """Load a MongoDocumentStore from a MongoDB URI."""
         mongo_kvstore = MongoDBKVStore.from_uri(uri, db_name)
-        return cls(mongo_kvstore, namespace)
+        return cls(
+            mongo_kvstore,
+            namespace,
+            node_collection_suffix,
+            ref_doc_collection_suffix,
+            metadata_collection_suffix,
+        )
 
     @classmethod
     def from_host_and_port(
@@ -43,7 +62,16 @@ class MongoDocumentStore(KVDocumentStore):
         port: int,
         db_name: Optional[str] = None,
         namespace: Optional[str] = None,
+        node_collection_suffix: Optional[str] = None,
+        ref_doc_collection_suffix: Optional[str] = None,
+        metadata_collection_suffix: Optional[str] = None,
     ) -> "MongoDocumentStore":
         """Load a MongoDocumentStore from a MongoDB host and port."""
         mongo_kvstore = MongoDBKVStore.from_host_and_port(host, port, db_name)
-        return cls(mongo_kvstore, namespace)
+        return cls(
+            mongo_kvstore,
+            namespace,
+            node_collection_suffix,
+            ref_doc_collection_suffix,
+            metadata_collection_suffix,
+        )

llama-index-integrations/storage/docstore/llama-index-storage-docstore-mongodb/pyproject.toml

@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-storage-docstore-mongodb"
 readme = "README.md"
-version = "0.1.2"
+version = "0.1.3"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"