Add Delete Function to Vector Store Drivers

griptape/drivers/vector/base_vector_store_driver.py

@@ -72,6 +72,10 @@ def upsert_text(
             **kwargs
         )
 
+    @abstractmethod
+    def delete_vector(self, vector_id: str):
+        ...
+
     @abstractmethod
     def upsert_vector(
         self,

griptape/drivers/vector/local_vector_store_driver.py

@@ -68,3 +68,6 @@ def query(
 
     def _namespaced_vector_id(self, vector_id: str, namespace: Optional[str]):
         return vector_id if namespace is None else f"{namespace}-{vector_id}"
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

griptape/drivers/vector/marqo_vector_store_driver.py

@@ -238,4 +238,7 @@ def upsert_vector(
             The ID of the vector that was added.
         """
 
-        raise Exception("not implemented")
+        raise NotImplementedError(f"{self.__class__.__name__} does not support upserting a vector.")
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

griptape/drivers/vector/mongodb_vector_store_driver.py

@@ -135,3 +135,12 @@ def query(
         ]
 
         return results
+
+    def delete_vector(self, vector_id: str):
+        """Deletes the vector from the collection."""
+        collection = self.get_collection()
+        collection.delete_one({"_id": vector_id})
+
+    def delete_collection(self):
+        """Deletes the current collection from the database."""
+        self.get_collection().drop()

griptape/drivers/vector/opensearch_vector_store_driver.py

@@ -193,3 +193,6 @@ def create_index(self, vector_dimension: int | None = None, settings_override: d
                 self.client.indices.create(index=self.index_name, body=mapping)
         except Exception as e:
             logging.error(f"Error while handling index: {e}")
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

griptape/drivers/vector/pgvector_vector_store_driver.py

@@ -174,3 +174,6 @@ class VectorModel(Base):
             meta = Column(JSON)
 
         return VectorModel
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

griptape/drivers/vector/pinecone_vector_store_driver.py

@@ -100,3 +100,6 @@ def create_index(self, name: str, **kwargs) -> None:
 
         pinecone = import_optional_dependency("pinecone")
         pinecone.create_index(**params)
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

griptape/drivers/vector/redis_vector_store_driver.py

@@ -172,3 +172,6 @@ def _generate_key(self, vector_id: str, namespace: str | None = None) -> str:
     def _get_doc_prefix(self, namespace: str | None = None) -> str:
         """Get the document prefix based on the provided namespace."""
         return f"{namespace}:" if namespace else ""
+
+    def delete_vector(self, vector_id: str):
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deletion.")

pyproject.toml

@@ -107,6 +107,7 @@ pytest = "~=7.1"
 pytest-mock = "*"
 mongomock = "*"
 
+
 twine = ">=4"
 moto = {extras = ["dynamodb"], version = "^4.1.8"}
 pytest-xdist = "^3.3.1"

tests/unit/drivers/vector/test_mongodb_vector_store_driver.py

@@ -1,6 +1,7 @@
 import pytest
 import mongomock
-from pymongo import MongoClient
+from unittest.mock import patch
+from pymongo.errors import OperationFailure
 from griptape.artifacts import TextArtifact
 from griptape.drivers import MongoDbAtlasVectorStoreDriver, BaseVectorStoreDriver
 from tests.mocks.mock_embedding_driver import MockEmbeddingDriver
@@ -64,3 +65,28 @@ def test_load_entries(self, driver):
         driver.upsert_vector(vector, vector_id=vector_id_str)  # ensure at least one entry exists
         results = list(driver.load_entries())
         assert results is not None and len(results) > 0
+
+    def test_delete(self, driver):
+        vector_id_str = "123"
+        vector = [0.5, 0.5, 0.5]
+        driver.upsert_vector(vector, vector_id=vector_id_str)  # ensure at least one entry exists
+        results = list(driver.load_entries())
+        assert results is not None and len(results) > 0
+
+        driver.delete_vector(vector_id_str)
+        results = list(driver.load_entries())
+        assert results is not None and len(results) == 0
+
+    def test_delete_collection(self, driver):
+        client = driver.client[driver.database_name]
+        # as mongomock follows MongoDB's convention of lazy collection creation
+        # we need to insert a document to ensure the collection exists
+        vector_id_str = "123"
+        vector = [0.5, 0.5, 0.5]
+        driver.upsert_vector(vector, vector_id=vector_id_str)
+        results = list(driver.load_entries())
+        assert results is not None and len(results) > 0
+        assert driver.get_collection().count_documents({}) == 1
+        assert len(client.list_collection_names()) == 1
+        driver.delete_collection()
+        assert len(client.list_collection_names()) == 0