feat: add mongo db integration

truefoundry · Dec 6, 2024 · 54baa0d · 54baa0d
1 parent db810a1
commit 54baa0d
Show file tree

Hide file tree

Showing 6 changed files with 347 additions and 98 deletions.
diff --git a/backend/modules/vector_db/__init__.py b/backend/modules/vector_db/__init__.py
@@ -1,5 +1,6 @@
 from backend.modules.vector_db.base import BaseVectorDB
-from backend.modules.vector_db.mongo import MongoVectorDB
+
+# from backend.modules.vector_db.mongo import MongoVectorDB
 from backend.modules.vector_db.qdrant import QdrantVectorDB
 
 # from backend.modules.vector_db.singlestore import SingleStoreVectorDB
@@ -8,7 +9,7 @@
 
 SUPPORTED_VECTOR_DBS = {
     "qdrant": QdrantVectorDB,
-    "mongo": MongoVectorDB,
+    # "mongo": MongoVectorDB,
     # "weaviate": WeaviateVectorDB,
     # "singlestore": SingleStoreVectorDB,
 }

diff --git a/backend/modules/vector_db/base.py b/backend/modules/vector_db/base.py
@@ -6,6 +6,7 @@
 from langchain.schema.vectorstore import VectorStore
 
 from backend.constants import DEFAULT_BATCH_SIZE_FOR_VECTOR_STORE
+from backend.logger import logger
 from backend.types import DataPointVector
 
 
@@ -88,4 +89,9 @@ def get_embedding_dimensions(self, embeddings: Embeddings) -> int:
         """
         Fetch embedding dimensions
         """
-        raise NotImplementedError()
+        # Calculate embedding size
+        logger.debug(f"[VectorDB] Embedding a dummy doc to get vector dimensions")
+        partial_embeddings = embeddings.embed_documents(["Initial document"])
+        vector_size = len(partial_embeddings[0])
+        logger.debug(f"Vector size: {vector_size}")
+        return vector_size