浏览代码

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.
google-labs-jules[bot] 1 周之前
父节点
当前提交
123dbf152e
共有 2 个文件被更改,包括 13 次插入1 次删除
  1. 4 0
      backend/open_webui/config.py
  2. 9 1
      backend/open_webui/retrieval/vector/dbs/milvus.py

+ 4 - 0
backend/open_webui/config.py

@@ -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)

+ 9 - 1
backend/open_webui/retrieval/vector/dbs/milvus.py

@@ -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.