diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 7e5c35a451..7eb43eb96f 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -2009,6 +2009,10 @@ MILVUS_METRIC_TYPE = os.environ.get("MILVUS_METRIC_TYPE", "COSINE") MILVUS_HNSW_M = int(os.environ.get("MILVUS_HNSW_M", "16")) MILVUS_HNSW_EFCONSTRUCTION = int(os.environ.get("MILVUS_HNSW_EFCONSTRUCTION", "100")) MILVUS_IVF_FLAT_NLIST = int(os.environ.get("MILVUS_IVF_FLAT_NLIST", "128")) +MILVUS_DISKANN_MAX_DEGREE = int(os.environ.get("MILVUS_DISKANN_MAX_DEGREE", "56")) +MILVUS_DISKANN_SEARCH_LIST_SIZE = int( + os.environ.get("MILVUS_DISKANN_SEARCH_LIST_SIZE", "100") +) # 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 53bb3bb72f..f7bd30cbd7 100644 --- a/backend/open_webui/retrieval/vector/dbs/milvus.py +++ b/backend/open_webui/retrieval/vector/dbs/milvus.py @@ -22,6 +22,8 @@ from open_webui.config import ( MILVUS_HNSW_M, MILVUS_HNSW_EFCONSTRUCTION, MILVUS_IVF_FLAT_NLIST, + MILVUS_DISKANN_MAX_DEGREE, + MILVUS_DISKANN_SEARCH_LIST_SIZE, ) from open_webui.env import SRC_LOG_LEVELS @@ -131,12 +133,18 @@ class MilvusClient(VectorDBBase): elif index_type == "IVF_FLAT": index_creation_params = {"nlist": MILVUS_IVF_FLAT_NLIST} log.info(f"IVF_FLAT params: {index_creation_params}") + elif index_type == "DISKANN": + index_creation_params = { + "max_degree": MILVUS_DISKANN_MAX_DEGREE, + "search_list_size": MILVUS_DISKANN_SEARCH_LIST_SIZE, + } + log.info(f"DISKANN params: {index_creation_params}") elif index_type in ["FLAT", "AUTOINDEX"]: log.info(f"Using {index_type} index with no specific build-time params.") else: log.warning( f"Unsupported MILVUS_INDEX_TYPE: '{index_type}'. " - f"Supported types: HNSW, IVF_FLAT, FLAT, AUTOINDEX. " + f"Supported types: HNSW, IVF_FLAT, DISKANN, FLAT, AUTOINDEX. " f"Milvus will use its default for the collection if this type is not directly supported for index creation." ) # For unsupported types, pass the type directly to Milvus; it might handle it or use a default.