Add a drop_documents() method to search index classes (#316)

Wrap key creation so users don't have to call private methods.
Co-authored-by: Tyler Hutcherson <[email protected]>

Author: abrookins

.github/workflows/test.yml

@@ -62,7 +62,7 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
           HF_HOME: ${{ github.workspace }}/hf_cache
-          OPENAI_API_KEY: ${{ secrets.OPENAI_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           GCP_LOCATION: ${{ secrets.GCP_LOCATION }}
           GCP_PROJECT_ID: ${{ secrets.GCP_PROJECT_ID }}
           COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
@@ -152,7 +152,7 @@ jobs:
         if: matrix.connection == 'plain' && matrix.redis-version == 'latest'
         env:
           HF_HOME: ${{ github.workspace }}/hf_cache
-          OPENAI_API_KEY: ${{ secrets.OPENAI_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           GCP_LOCATION: ${{ secrets.GCP_LOCATION }}
           GCP_PROJECT_ID: ${{ secrets.GCP_PROJECT_ID }}
           COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}

docs/overview/cli.ipynb

@@ -24,7 +24,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\u001b[32m16:19:10\u001b[0m \u001b[34m[RedisVL]\u001b[0m \u001b[1;30mINFO\u001b[0m   RedisVL version 0.4.0\n"
+      "\u001b[32m09:58:03\u001b[0m \u001b[34m[RedisVL]\u001b[0m \u001b[1;30mINFO\u001b[0m   RedisVL version 0.5.2\n"
      ]
     }
    ],

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redisvl"
-version = "0.5.1"
+version = "0.5.2"
 description = "Python client library and CLI for using Redis as a vector database"
 authors = ["Redis Inc. <[email protected]>"]
 license = "MIT"
@@ -31,7 +31,7 @@ redis = "^5.0"
 pydantic = "^2"
 tenacity = ">=8.2.2"
 tabulate = "^0.9.0"
-ml-dtypes = "^0.4.0"
+ml-dtypes = ">=0.4.0,<1.0.0"
 python-ulid = "^3.0.0"
 nltk = { version = "^3.8.1", optional = true }
 jsonpath-ng = "^1.5.0"

pyproject.toml

@@ -592,6 +592,27 @@ def drop_keys(self, keys: Union[str, List[str]]) -> int:
         else:
             return self._redis_client.delete(keys)  # type: ignore
 
+    def drop_documents(self, ids: Union[str, List[str]]) -> int:
+        """Remove documents from the index by their document IDs.
+
+        This method converts document IDs to Redis keys automatically by applying
+        the index's key prefix and separator configuration.
+
+        Args:
+            ids (Union[str, List[str]]): The document ID or IDs to remove from the index.
+
+        Returns:
+            int: Count of documents deleted from Redis.
+        """
+        if isinstance(ids, list):
+            if not ids:
+                return 0
+            keys = [self.key(id) for id in ids]
+            return self._redis_client.delete(*keys)  # type: ignore
+        else:
+            key = self.key(ids)
+            return self._redis_client.delete(key)  # type: ignore
+
     def expire_keys(
         self, keys: Union[str, List[str]], ttl: int
     ) -> Union[int, List[int]]:
@@ -1236,6 +1257,28 @@ async def drop_keys(self, keys: Union[str, List[str]]) -> int:
         else:
             return await client.delete(keys)
 
+    async def drop_documents(self, ids: Union[str, List[str]]) -> int:
+        """Remove documents from the index by their document IDs.
+
+        This method converts document IDs to Redis keys automatically by applying
+        the index's key prefix and separator configuration.
+
+        Args:
+            ids (Union[str, List[str]]): The document ID or IDs to remove from the index.
+
+        Returns:
+            int: Count of documents deleted from Redis.
+        """
+        client = await self._get_client()
+        if isinstance(ids, list):
+            if not ids:
+                return 0
+            keys = [self.key(id) for id in ids]
+            return await client.delete(*keys)
+        else:
+            key = self.key(ids)
+            return await client.delete(key)
+
     async def expire_keys(
         self, keys: Union[str, List[str]], ttl: int
     ) -> Union[int, List[int]]:
@@ -1356,7 +1399,7 @@ async def fetch(self, id: str) -> Optional[Dict[str, Any]]:
     async def _aggregate(
         self, aggregation_query: AggregationQuery
     ) -> List[Dict[str, Any]]:
-        """Execute an aggretation query and processes the results."""
+        """Execute an aggregation query and processes the results."""
         results = await self.aggregate(
             aggregation_query, query_params=aggregation_query.params  # type: ignore[attr-defined]
         )

redisvl/index/index.py

@@ -1 +1 @@
-__version__ = "0.5.1"
+__version__ = "0.5.2"

redisvl/version.py

@@ -246,6 +246,37 @@ async def test_search_index_drop_keys(async_index):
     assert await async_index.exists()
 
 
+@pytest.mark.asyncio
+async def test_search_index_drop_documents(async_index):
+    await async_index.create(overwrite=True, drop=True)
+    data = [
+        {"id": "1", "test": "foo"},
+        {"id": "2", "test": "bar"},
+        {"id": "3", "test": "baz"},
+    ]
+    await async_index.load(data, id_field="id")
+
+    # Test dropping a single document by ID
+    dropped = await async_index.drop_documents("1")
+    assert dropped == 1
+    assert not await async_index.fetch("1")
+    assert await async_index.fetch("2") is not None
+    assert await async_index.fetch("3") is not None
+
+    # Test dropping multiple documents by ID
+    dropped = await async_index.drop_documents(["2", "3"])
+    assert dropped == 2
+    assert not await async_index.fetch("2")
+    assert not await async_index.fetch("3")
+
+    # Test dropping with an empty list
+    dropped = await async_index.drop_documents([])
+    assert dropped == 0
+
+    # Ensure the index still exists
+    assert await async_index.exists()
+
+
 @pytest.mark.asyncio
 async def test_search_index_load_and_fetch(async_index):
     await async_index.create(overwrite=True, drop=True)

tests/integration/test_async_search_index.py

@@ -240,6 +240,36 @@ def test_search_index_drop_keys(index):
     assert index.exists()
 
 
+def test_search_index_drop_documents(index):
+    index.create(overwrite=True, drop=True)
+    data = [
+        {"id": "1", "test": "foo"},
+        {"id": "2", "test": "bar"},
+        {"id": "3", "test": "baz"},
+    ]
+    index.load(data, id_field="id")
+
+    # Test dropping a single document by ID
+    dropped = index.drop_documents("1")
+    assert dropped == 1
+    assert not index.fetch("1")
+    assert index.fetch("2") is not None
+    assert index.fetch("3") is not None
+
+    # Test dropping multiple documents by ID
+    dropped = index.drop_documents(["2", "3"])
+    assert dropped == 2
+    assert not index.fetch("2")
+    assert not index.fetch("3")
+
+    # Test dropping with an empty list
+    dropped = index.drop_documents([])
+    assert dropped == 0
+
+    # Ensure the index still exists
+    assert index.exists()
+
+
 def test_search_index_load_and_fetch(index):
     index.create(overwrite=True, drop=True)
     data = [{"id": "1", "test": "foo"}]