UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-02-21 12:29:29 +00:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

refac: vector db clients

Browse Source
This commit is contained in:
Timothy J. Baek 2024-09-12 02:00:31 -04:00
parent 4775fe43d8
commit 8be6e16513
3 changed files with 9 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
backend/open_webui/apps/rag/main.py
View File

@ -997,7 +997,7 @@ def store_docs_in_vector_db(
try: try:
if overwrite: if overwrite:
if collection_name in VECTOR_DB_CLIENT.list_collections(): if VECTOR_DB_CLIENT.has_collection(collection_name=collection_name):
log.info(f"deleting existing collection {collection_name}") log.info(f"deleting existing collection {collection_name}")
VECTOR_DB_CLIENT.delete_collection(collection_name=collection_name) VECTOR_DB_CLIENT.delete_collection(collection_name=collection_name)

6
backend/open_webui/apps/rag/vector/dbs/chroma.py
View File

@ -36,10 +36,10 @@ class ChromaClient:
database=CHROMA_DATABASE, database=CHROMA_DATABASE,
) )
def list_collections(self) -> list[str]: def has_collection(self, collection_name: str) -> bool:
# List all the collections in the database. # Check if the collection exists based on the collection name.
collections = self.client.list_collections() collections = self.client.list_collections()
return [collection.name for collection in collections] return collection_name in [collection.name for collection in collections]
def delete_collection(self, collection_name: str): def delete_collection(self, collection_name: str):
# Delete the collection based on the collection name. # Delete the collection based on the collection name.

12
backend/open_webui/apps/rag/vector/dbs/milvus.py
View File

@ -80,13 +80,11 @@ class MilvusClient:
index_params=index_params, index_params=index_params,
) )
def list_collections(self) -> list[str]: def has_collection(self, collection_name: str) -> bool:
# List all the collections in the database. # Check if the collection exists based on the collection name.
return [ return self.client.has_collection(
collection[len(self.collection_prefix) :] collection_name=f"{self.collection_prefix}_{collection_name}"
for collection in self.client.list_collections() )
if collection.startswith(self.collection_prefix)
]
def delete_collection(self, collection_name: str): def delete_collection(self, collection_name: str):
# Delete the collection based on the collection name. # Delete the collection based on the collection name.