UniqAI
/
open-webui
mirror of https://github.com/open-webui/open-webui
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #5364 from open-webui/multiple-vector-dbs 

refac: vector db clients
This commit is contained in:
Timothy Jaeryang Baek 2024-09-12 07:00:55 +01:00 committed by GitHub
parent c7fc17da69 8be6e16513
commit 675403d26d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
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:
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}")
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,
)
def list_collections(self) -> list[str]:
# List all the collections in the database.
def has_collection(self, collection_name: str) -> bool:
# Check if the collection exists based on the collection name.
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):
# 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,
)
def list_collections(self) -> list[str]:
# List all the collections in the database.
return [
collection[len(self.collection_prefix) :]
for collection in self.client.list_collections()
if collection.startswith(self.collection_prefix)
]
def has_collection(self, collection_name: str) -> bool:
# Check if the collection exists based on the collection name.
return self.client.has_collection(
collection_name=f"{self.collection_prefix}_{collection_name}"
)
def delete_collection(self, collection_name: str):
# Delete the collection based on the collection name.