From 123dbf152e4651b3789cdd560873986963f732ac Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Fri, 26 Sep 2025 06:54:06 +0000
Subject: [PATCH] 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.
---
 backend/open_webui/config.py                      |  4 ++++
 backend/open_webui/retrieval/vector/dbs/milvus.py | 10 +++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index ca090efa22..d3ec3b8c3d 100644
--- a/backend/open_webui/config.py
+++ b/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_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 059ea43cc0..ae1485cfc2 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.