Merge pull request #9191 from NYU-ITS/milvus_auth_token

[feat] Milvus: new optional config var, MILVUS_TOKEN
This commit is contained in:
Timothy Jaeryang Baek 2025-02-01 21:04:30 -08:00 committed by GitHub
commit 95c94d8a85
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 6 additions and 1 deletions

View File

@ -1346,6 +1346,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_URI = os.environ.get("MILVUS_URI", f"{DATA_DIR}/vector_db/milvus.db")
MILVUS_DB = os.environ.get("MILVUS_DB", "default") MILVUS_DB = os.environ.get("MILVUS_DB", "default")
MILVUS_TOKEN = os.environ.get("MILVUS_TOKEN", None)
# Qdrant # Qdrant
QDRANT_URI = os.environ.get("QDRANT_URI", None) QDRANT_URI = os.environ.get("QDRANT_URI", None)

View File

@ -8,13 +8,17 @@ from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
from open_webui.config import ( from open_webui.config import (
MILVUS_URI, MILVUS_URI,
MILVUS_DB, MILVUS_DB,
MILVUS_TOKEN,
) )
class MilvusClient: class MilvusClient:
def __init__(self): def __init__(self):
self.collection_prefix = "open_webui" 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: def _result_to_get_result(self, result) -> GetResult:
ids = [] ids = []