Browse Source

enh: vector db delete filter support

Timothy J. Baek 7 months ago
parent
commit
325ca98773

+ 10 - 2
backend/open_webui/apps/retrieval/vector/dbs/chroma.py

@@ -111,11 +111,19 @@ class ChromaClient:
             ids=ids, documents=documents, embeddings=embeddings, metadatas=metadatas
         )
 
-    def delete(self, collection_name: str, ids: list[str]):
+    def delete(
+        self,
+        collection_name: str,
+        ids: Optional[list[str]] = None,
+        filter: Optional[dict] = None,
+    ):
         # Delete the items from the collection based on the ids.
         collection = self.client.get_collection(name=collection_name)
         if collection:
-            collection.delete(ids=ids)
+            if ids:
+                collection.delete(ids=ids)
+            elif filter:
+                collection.delete(where=filter)
 
     def reset(self):
         # Resets the database. This will delete all collections and item entries.

+ 24 - 5
backend/open_webui/apps/retrieval/vector/dbs/milvus.py

@@ -187,13 +187,32 @@ class MilvusClient:
             ],
         )
 
-    def delete(self, collection_name: str, ids: list[str]):
+    def delete(
+        self,
+        collection_name: str,
+        ids: Optional[list[str]] = None,
+        filter: Optional[dict] = None,
+    ):
         # Delete the items from the collection based on the ids.
 
-        return self.client.delete(
-            collection_name=f"{self.collection_prefix}_{collection_name}",
-            ids=ids,
-        )
+        if ids:
+            return self.client.delete(
+                collection_name=f"{self.collection_prefix}_{collection_name}",
+                ids=ids,
+            )
+        elif filter:
+            # Convert the filter dictionary to a string using JSON_CONTAINS.
+            filter_string = " && ".join(
+                [
+                    f"JSON_CONTAINS(metadata[{key}], '{[value] if isinstance(value, str) else value}')"
+                    for key, value in filter.items()
+                ]
+            )
+
+            return self.client.delete(
+                collection_name=f"{self.collection_prefix}_{collection_name}",
+                filter=filter_string,
+            )
 
     def reset(self):
         # Resets the database. This will delete all collections and item entries.