Skip to content

Commit

Permalink
feat: Add support for LanceDB as Vector DB provider
Browse files Browse the repository at this point in the history
  • Loading branch information
sp6370 committed Mar 26, 2024
1 parent 1fdf426 commit 0d085e0
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 1 deletion.
2 changes: 1 addition & 1 deletion r2r/core/providers/vector_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
54 changes: 54 additions & 0 deletions r2r/vector_dbs/lancedb/base.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
import logging
import os
from typing import Optional, Union

from r2r.core import VectorDBProvider, VectorEntry, VectorSearchResult

logger = logging.getLogger(__name__)


class LanceDB(VectorDBProvider):
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(db_path=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 # TODO ADD Dependency
except ImportError:
raise ValueError(
f"Error, `pyarrow` is not installed. Please install it using `pip install pyarrow`."
)
try:
result = self.client.create_table(
name=f"{collection_name}", on_bad_vectors="error", schema=[]
)
except Exception:
# TODO - Handle more appropriately - create collection fails when it already exists
# https://lancedb.github.io/lancedb/python/python/#lancedb.db.DBConnection.create_table
print(Exception)
pass

0 comments on commit 0d085e0

Please sign in to comment.