Deeplake tests, typing

src/grag/components/vectordb/base.py

@@ -1,13 +1,23 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Tuple, Union
 
 from langchain_community.vectorstores.utils import filter_complex_metadata
 from langchain_core.documents import Document
 
 
 class VectorDB(ABC):
+
+    @abstractmethod
+    def __len__(self) -> int:
+        """Number of chunks in the vector database."""
+        ...
+
+    @abstractmethod
+    def delete(self) -> None:
+        """Delete all chunks in the vector database."""
+
     @abstractmethod
-    def add_docs(self, docs: List[Document], verbose: bool = True):
+    def add_docs(self, docs: List[Document], verbose: bool = True) -> None:
         """Adds documents to the vector database.
         
         Args:
@@ -20,7 +30,7 @@ def add_docs(self, docs: List[Document], verbose: bool = True):
         ...
 
     @abstractmethod
-    async def aadd_docs(self, docs: List[Document], verbose: bool = True):
+    async def aadd_docs(self, docs: List[Document], verbose: bool = True) -> None:
         """Adds documents to the vector database (asynchronous).
 
         Args:
@@ -33,7 +43,8 @@ async def aadd_docs(self, docs: List[Document], verbose: bool = True):
         ...
 
     @abstractmethod
-    def get_chunk(self, query: str, with_score: bool = False, top_k: int = None):
+    def get_chunk(self, query: str, with_score: bool = False, top_k: int = None) -> Union[
+        List[Document], List[Tuple[Document, float]]]:
         """Returns the most similar chunks from the vector database.
 
         Args:
@@ -47,7 +58,8 @@ def get_chunk(self, query: str, with_score: bool = False, top_k: int = None):
         ...
 
     @abstractmethod
-    async def aget_chunk(self, query: str, with_score: bool = False, top_k: int = None):
+    async def aget_chunk(self, query: str, with_score: bool = False, top_k: int = None) -> Union[
+        List[Document], List[Tuple[Document, float]]]:
         """Returns the most similar chunks from the vector database. (asynchronous)
 
         Args:
@@ -60,5 +72,5 @@ async def aget_chunk(self, query: str, with_score: bool = False, top_k: int = No
         """
         ...
 
-    def _filter_metadata(self, docs: List[Document]):
+    def _filter_metadata(self, docs: List[Document]) -> List[Document]:
         return filter_complex_metadata(docs, allowed_types=self.allowed_metadata_types)

src/grag/components/vectordb/chroma_client.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple, Union
 
 import chromadb
 from grag.components.embedding import Embedding
@@ -72,7 +72,21 @@ def __init__(
         )
         self.allowed_metadata_types = (str, int, float, bool)
 
-    def test_connection(self, verbose=True):
+    def __len__(self) -> int:
+        return self.collection.count()
+
+    def delete(self) -> None:
+        self.client.delete_collection(self.collection_name)
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name
+        )
+        self.langchain_client = Chroma(
+            client=self.client,
+            collection_name=self.collection_name,
+            embedding_function=self.embedding_function,
+        )
+
+    def test_connection(self, verbose=True) -> int:
         """Tests connection with Chroma Vectorstore
 
         Args:
@@ -89,7 +103,7 @@ def test_connection(self, verbose=True):
                 print(f"Connection to {self.host}/{self.port} is not alive !!")
         return response
 
-    def add_docs(self, docs: List[Document], verbose=True):
+    def add_docs(self, docs: List[Document], verbose=True) -> None:
         """Adds documents to chroma vectorstore
 
         Args:
@@ -105,7 +119,7 @@ def add_docs(self, docs: List[Document], verbose=True):
         ):
             _id = self.langchain_client.add_documents([doc])
 
-    async def aadd_docs(self, docs: List[Document], verbose=True):
+    async def aadd_docs(self, docs: List[Document], verbose=True) -> None:
         """Asynchronously adds documents to chroma vectorstore
 
         Args:
@@ -127,7 +141,8 @@ async def aadd_docs(self, docs: List[Document], verbose=True):
             for doc in docs:
                 await self.langchain_client.aadd_documents([doc])
 
-    def get_chunk(self, query: str, with_score: bool = False, top_k: int = None):
+    def get_chunk(self, query: str, with_score: bool = False, top_k: int = None) -> Union[
+        List[Document], List[Tuple[Document, float]]]:
         """Returns the most similar chunks from the chroma database.
 
         Args:
@@ -148,7 +163,8 @@ def get_chunk(self, query: str, with_score: bool = False, top_k: int = None):
                 query=query, k=top_k if top_k else 1
             )
 
-    async def aget_chunk(self, query: str, with_score=False, top_k=None):
+    async def aget_chunk(self, query: str, with_score=False, top_k=None) -> Union[
+        List[Document], List[Tuple[Document, float]]]:
         """Returns the most (cosine) similar chunks from the vector database, asynchronously.
 
         Args:

src/grag/components/vectordb/deeplake_client.py

@@ -1,7 +1,6 @@
 from pathlib import Path
-from typing import List, Union
+from typing import List, Tuple, Union
 
-from deeplake.core.vectorstore import VectorStore
 from grag.components.embedding import Embedding
 from grag.components.utils import get_config
 from grag.components.vectordb.base import VectorDB
@@ -34,24 +33,36 @@ class DeepLakeClient(VectorDB):
     """
 
     def __init__(self,
-                 store_path: Union[str, Path],
-                 embedding_model: str,
-                 embedding_type: str,
+                 collection_name: str = deeplake_conf["collection_name"],
+                 store_path: Union[str, Path] = deeplake_conf["store_path"],
+                 embedding_type: str = deeplake_conf["embedding_type"],
+                 embedding_model: str = deeplake_conf["embedding_model"],
+                 read_only: bool = False
                  ):
         self.store_path = Path(store_path)
+        self.collection_name = collection_name
+        self.read_only = read_only
         self.embedding_type: str = embedding_type
         self.embedding_model: str = embedding_model
 
         self.embedding_function = Embedding(
             embedding_model=self.embedding_model, embedding_type=self.embedding_type
         ).embedding_function
 
-        self.client = VectorStore(path=self.store_path)
-        self.langchain_client = DeepLake(path=self.store_path,
-                                         embedding=self.embedding_function)
+        # self.client = VectorStore(path=self.store_path / self.collection_name)
+        self.langchain_client = DeepLake(dataset_path=str(self.store_path / self.collection_name),
+                                         embedding=self.embedding_function,
+                                         read_only=self.read_only)
+        self.client = self.langchain_client.vectorstore
         self.allowed_metadata_types = (str, int, float, bool)
 
-    def add_docs(self, docs: List[Document], verbose=True):
+    def __len__(self) -> int:
+        return self.client.__len__()
+
+    def delete(self) -> None:
+        self.client.delete(delete_all=True)
+
+    def add_docs(self, docs: List[Document], verbose=True) -> None:
         """Adds documents to deeplake vectorstore
     
         Args:
@@ -65,9 +76,9 @@ def add_docs(self, docs: List[Document], verbose=True):
         for doc in (
                 tqdm(docs, desc=f"Adding to {self.collection_name}:") if verbose else docs
         ):
-            _id = self.langchain_chroma.add_documents([doc])
+            _id = self.langchain_client.add_documents([doc])
 
-    async def aadd_docs(self, docs: List[Document], verbose=True):
+    async def aadd_docs(self, docs: List[Document], verbose=True) -> None:
         """Asynchronously adds documents to chroma vectorstore
     
         Args:
@@ -84,12 +95,13 @@ async def aadd_docs(self, docs: List[Document], verbose=True):
                     desc=f"Adding documents to {self.collection_name}",
                     total=len(docs),
             ):
-                await self.langchain_deeplake.aadd_documents([doc])
+                await self.langchain_client.aadd_documents([doc])
         else:
             for doc in docs:
-                await self.langchain_deeplake.aadd_documents([doc])
+                await self.langchain_client.aadd_documents([doc])
 
-    def get_chunk(self, query: str, with_score: bool = False, top_k: int = None):
+    def get_chunk(self, query: str, with_score: bool = False, top_k: int = None) -> Union[
+        List[Document], List[Tuple[Document, float]]]:
         """Returns the most similar chunks from the deeplake database.
 
         Args:
@@ -102,15 +114,16 @@ def get_chunk(self, query: str, with_score: bool = False, top_k: int = None):
 
         """
         if with_score:
-            return self.langchain_client.similarity_search_with_relevance_scores(
+            return self.langchain_client.similarity_search_with_score(
                 query=query, k=top_k if top_k else 1
             )
         else:
             return self.langchain_client.similarity_search(
                 query=query, k=top_k if top_k else 1
             )
 
-    async def aget_chunk(self, query: str, with_score=False, top_k=None):
+    async def aget_chunk(self, query: str, with_score=False, top_k=None) -> Union[
+        List[Document], List[Tuple[Document, float]]]:
         """Returns the most similar chunks from the deeplake database, asynchronously.
 
         Args:
@@ -123,7 +136,7 @@ async def aget_chunk(self, query: str, with_score=False, top_k=None):
 
         """
         if with_score:
-            return await self.langchain_client.asimilarity_search_with_relevance_scores(
+            return await self.langchain_client.asimilarity_search_with_score(
                 query=query, k=top_k if top_k else 1
             )
         else:

src/tests/components/vectordb/chroma_client_test.py

@@ -47,13 +47,11 @@ def test_chroma_add_docs():
     thunder rolled and boomed, like the Colorado in flood.""",
     ]
     chroma_client = ChromaClient(collection_name="test")
-    if chroma_client.collection.count() > 0:
-        chroma_client.client.delete_collection("test")
-    chroma_client = ChromaClient(collection_name="test")
+    if len(chroma_client) > 0:
+        chroma_client.delete()
     docs = [Document(page_content=doc) for doc in docs]
     chroma_client.add_docs(docs)
-    collection_count = chroma_client.collection.count()
-    assert collection_count == len(docs)
+    assert len(chroma_client) == len(docs)
 
 
 def test_chroma_aadd_docs():
@@ -92,13 +90,12 @@ def test_chroma_aadd_docs():
     thunder rolled and boomed, like the Colorado in flood.""",
     ]
     chroma_client = ChromaClient(collection_name="test")
-    if chroma_client.collection.count() > 0:
-        chroma_client.client.delete_collection("test")
-    chroma_client = ChromaClient(collection_name="test")
+    if len(chroma_client) > 0:
+        chroma_client.delete()
     docs = [Document(page_content=doc) for doc in docs]
     loop = asyncio.get_event_loop()
     loop.run_until_complete(chroma_client.aadd_docs(docs))
-    assert chroma_client.collection.count() == len(docs)
+    assert len(chroma_client) == len(docs)
 
 
 chrome_get_chunk_params = [(1, False), (1, True), (2, False), (2, True)]