[WIP] feat: add support for LanceDB as vector db provider

pyproject.toml

@@ -30,18 +30,20 @@ uvicorn = "^0.27.0.post1"
 vecs = "^0.4.0"
 
 # Optional dependencies
-bs4 = {version = "^0.0.2", optional = true}
-pypdf = {version = "^4.0.2", optional = true}
-tiktoken = {version = "^0.5.2", optional = true}
-datasets = {version = "^2.16.1", optional = true}
-qdrant_client = {version = "^1.7.0", optional = true}
-psycopg2-binary = {version = "^2.9.9", optional = true}
-numpy = {version = "^1.26.4", optional = true}
-scikit-learn = {version = "^1.4.1.post1", optional = true}
-sentry-sdk = {version = "^1.40.4", optional = true}
-deepeval = {version ="^0.20.88", optional = true}
-parea-ai = {version = "^0.2.86", optional = true}
-ionic-api-sdk = {version = "0.9.3", optional = true}
+bs4 = { version = "^0.0.2", optional = true }
+pypdf = { version = "^4.0.2", optional = true }
+tiktoken = { version = "^0.5.2", optional = true }
+datasets = { version = "^2.16.1", optional = true }
+qdrant_client = { version = "^1.7.0", optional = true }
+psycopg2-binary = { version = "^2.9.9", optional = true }
+numpy = { version = "^1.26.4", optional = true }
+scikit-learn = { version = "^1.4.1.post1", optional = true }
+sentry-sdk = { version = "^1.40.4", optional = true }
+deepeval = { version = "^0.20.88", optional = true }
+parea-ai = { version = "^0.2.86", optional = true }
+ionic-api-sdk = { version = "0.9.3", optional = true }
+lancedb = { version = "^0.6.5", optional = true }
+pyarrow = { version = "^15.0.2", optional = true }
 
 [tool.poetry.extras]
 parsing = ["bs4", "pypdf"]
@@ -53,7 +55,8 @@ local_vectordb = ["numpy", "scikit-learn"]
 monitoring = ["sentry-sdk"]
 eval = ["deepeval", "parea-ai"]
 ionic = ["ionic-api-sdk"]
-all = ["bs4", "pypdf", "tiktoken", "datasets", "qdrant_client", "psycopg2-binary", "numpy", "scikit-learn", "sentry-sdk", "protobuf", "deepeval", "parea-ai", "ionic"]
+lancedb = ["lancedb"]
+all = ["bs4", "pypdf", "tiktoken", "datasets", "qdrant_client", "psycopg2-binary", "numpy", "scikit-learn", "sentry-sdk", "protobuf", "deepeval", "parea-ai", "ionic", "lancedb"]
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.3.0"

r2r/core/providers/vector_db.py

@@ -64,7 +64,7 @@ def to_dict(self) -> dict:
 
 
 class VectorDBProvider(ABC):
-    supported_providers = ["local", "pgvector", "qdrant"]
+    supported_providers = ["local", "pgvector", "qdrant", "lancedb"]
 
     def __init__(self, provider: str):
         if provider not in VectorDBProvider.supported_providers:

r2r/vector_dbs/lancedb/base.py

@@ -0,0 +1,89 @@
+import logging
+from typing import Optional
+
+from r2r.core import VectorDBProvider, VectorEntry
+
+logger = logging.getLogger(__name__)
+
+
+class LanceDB(VectorDBProvider):
+    # TODO enable LanceDB provider to support lanceDB Cloud
+    def __init__(
+            self, provider: str = "lancedb", db_path: Optional[str] = None
+    ) -> None:
+        logger.info("Initializing `LanceDB` to store and retrieve embeddings.")
+
+        super().__init__(provider)
+
+        if provider != "lancedb":
+            raise ValueError(
+                "LanceDB must be initialized with provider `lancedb`."
+            )
+
+        try:
+            import lancedb
+        except ImportError:
+            raise ValueError(
+                f"Error, `lancedb` is not installed. Please install it using `pip install lancedb`."
+            )
+
+        self.db_path = db_path
+
+        try:
+            self.client = lancedb.connect(uri=self.db_path)
+        except Exception as e:
+            raise ValueError(
+                f"Error {e} occurred while attempting to connect to the lancedb provider."
+            )
+        self.collection_name: Optional[str] = None
+
+    def initialize_collection(
+            self, collection_name: str, dimension: int
+    ) -> None:
+        self.collection_name = collection_name
+
+        try:
+            import pyarrow
+        except ImportError:
+            raise ValueError(
+                f"Error, `pyarrow` is not installed. Please install it using `pip install pyarrow`."
+            )
+
+        table_schema = pyarrow.schema(
+            [
+                pyarrow.field("id", pyarrow.string()),
+                pyarrow.field(
+                    "vector", pyarrow.list_(pyarrow.float32(), dimension)
+                ),
+                # TODO Handle storing metadata
+            ]
+        )
+
+        try:
+            self.client.create_table(
+                name=f"{collection_name}",
+                on_bad_vectors="error",
+                schema=table_schema,
+            )
+        except Exception as e:
+            # TODO - Handle more appropriately - create collection fails when it already exists
+            pass
+
+    def copy(self, entry: VectorEntry, commit=True) -> None:
+        raise NotImplementedError(
+            "LanceDB does not support the `copy` method."
+        )
+
+    def upsert(self, entry: VectorEntry, commit=True) -> None:
+        if self.collection_name is None:
+            raise ValueError(
+                "Please call `initialize_collection` before attempting to run `upsert`."
+            )
+        self.client.open_table(self.collection_name).add(
+            {
+                "vector": entry.vector,
+                "id": entry.id,
+                # TODO ADD metadata storage
+            },
+            mode="overwrite",
+        )