From 7b31c75271bdd8f5b575e49fa66bacadbe9ae6b0 Mon Sep 17 00:00:00 2001 From: Sajid Ali Date: Fri, 31 Jan 2025 17:00:38 -0500 Subject: [PATCH] Milvus: new optional config var, MILVUS_TOKEN modified: backend/open_webui/config.py modified: backend/open_webui/retrieval/vector/dbs/milvus.py --- backend/open_webui/config.py | 1 + backend/open_webui/retrieval/vector/dbs/milvus.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index d60f3d436..3bfb403a7 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -1325,6 +1325,7 @@ CHROMA_HTTP_SSL = os.environ.get("CHROMA_HTTP_SSL", "false").lower() == "true" MILVUS_URI = os.environ.get("MILVUS_URI", f"{DATA_DIR}/vector_db/milvus.db") MILVUS_DB = os.environ.get("MILVUS_DB", "default") +MILVUS_TOKEN = os.environ.get("MILVUS_TOKEN", None) # Qdrant QDRANT_URI = os.environ.get("QDRANT_URI", None) diff --git a/backend/open_webui/retrieval/vector/dbs/milvus.py b/backend/open_webui/retrieval/vector/dbs/milvus.py index bdfa16eb6..43c3f3d1a 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus.py @@ -8,13 +8,17 @@ from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult from open_webui.config import ( MILVUS_URI, MILVUS_DB, + MILVUS_TOKEN, ) class MilvusClient: def __init__(self): self.collection_prefix = "open_webui" - self.client = Client(uri=MILVUS_URI, database=MILVUS_DB) + if MILVUS_TOKEN is None: + self.client = Client(uri=MILVUS_URI, database=MILVUS_DB) + else: + self.client = Client(uri=MILVUS_URI, database=MILVUS_DB, token=MILVUS_TOKEN) def _result_to_get_result(self, result) -> GetResult: ids = []