feat: Add DISKANN index type support for Milvus

This commit introduces support for the DISKANN index type in the Milvus vector database integration.

Changes include:
- Added `MILVUS_DISKANN_MAX_DEGREE` and `MILVUS_DISKANN_SEARCH_LIST_SIZE` configuration variables.
- Updated the Milvus client to recognize and configure the DISKANN index type during collection creation.
This commit is contained in:
google-labs-jules[bot] 2025-09-26 06:54:06 +00:00
parent b55a38ee97
commit 123dbf152e
2 changed files with 13 additions and 1 deletions

View file

@ -2007,6 +2007,10 @@ MILVUS_METRIC_TYPE = os.environ.get("MILVUS_METRIC_TYPE", "COSINE")
MILVUS_HNSW_M = int(os.environ.get("MILVUS_HNSW_M", "16")) MILVUS_HNSW_M = int(os.environ.get("MILVUS_HNSW_M", "16"))
MILVUS_HNSW_EFCONSTRUCTION = int(os.environ.get("MILVUS_HNSW_EFCONSTRUCTION", "100")) 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_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
QDRANT_URI = os.environ.get("QDRANT_URI", None) QDRANT_URI = os.environ.get("QDRANT_URI", None)

View file

@ -22,6 +22,8 @@ from open_webui.config import (
MILVUS_HNSW_M, MILVUS_HNSW_M,
MILVUS_HNSW_EFCONSTRUCTION, MILVUS_HNSW_EFCONSTRUCTION,
MILVUS_IVF_FLAT_NLIST, MILVUS_IVF_FLAT_NLIST,
MILVUS_DISKANN_MAX_DEGREE,
MILVUS_DISKANN_SEARCH_LIST_SIZE,
) )
from open_webui.env import SRC_LOG_LEVELS from open_webui.env import SRC_LOG_LEVELS
@ -131,12 +133,18 @@ class MilvusClient(VectorDBBase):
elif index_type == "IVF_FLAT": elif index_type == "IVF_FLAT":
index_creation_params = {"nlist": MILVUS_IVF_FLAT_NLIST} index_creation_params = {"nlist": MILVUS_IVF_FLAT_NLIST}
log.info(f"IVF_FLAT params: {index_creation_params}") 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"]: elif index_type in ["FLAT", "AUTOINDEX"]:
log.info(f"Using {index_type} index with no specific build-time params.") log.info(f"Using {index_type} index with no specific build-time params.")
else: else:
log.warning( log.warning(
f"Unsupported MILVUS_INDEX_TYPE: '{index_type}'. " 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." 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. # For unsupported types, pass the type directly to Milvus; it might handle it or use a default.