UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-06 16:59:42 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

enh: vector db delete filter support

Browse Source
This commit is contained in:
Timothy J. Baek 2024-10-03 06:43:50 -07:00
parent 351b1dbf31
commit 325ca98773
2 changed files with 34 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
backend/open_webui/apps/retrieval/vector/dbs/chroma.py
View File

@ -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.

29
backend/open_webui/apps/retrieval/vector/dbs/milvus.py
View File

@ -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.