OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac: streamline Milvus index type handling using configuration options

Browse Source
This commit is contained in:
Matt Harrison 2025-05-07 23:39:56 -04:00
parent 2b8be86f23
commit 731251d11a
1 changed files with 25 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

78
backend/open_webui/retrieval/vector/dbs/milvus.py
View File

@ -14,6 +14,11 @@ from open_webui.config import (
MILVUS_URI, MILVUS_URI,
MILVUS_DB, MILVUS_DB,
MILVUS_TOKEN, MILVUS_TOKEN,
MILVUS_INDEX_TYPE,
MILVUS_METRIC_TYPE,
MILVUS_HNSW_M,
MILVUS_HNSW_EFCONSTRUCTION,
MILVUS_IVF_FLAT_NLIST,
) )
from open_webui.env import SRC_LOG_LEVELS from open_webui.env import SRC_LOG_LEVELS
@ -106,66 +111,33 @@ class MilvusClient(VectorDBBase):
index_params = self.client.prepare_index_params() index_params = self.client.prepare_index_params()
# Get index type from environment variable. # Use configurations from config.py
# Milvus standalone (local mode) supports: FLAT, IVF_FLAT, AUTOINDEX. index_type = MILVUS_INDEX_TYPE.upper()
# HNSW is often preferred for performance but may require a clustered Milvus setup. metric_type = MILVUS_METRIC_TYPE.upper()
# Defaulting to AUTOINDEX for broader compatibility, especially with Milvus standalone.
default_index_type = "AUTOINDEX" log.info(f"Using Milvus index type: {index_type}, metric type: {metric_type}")
milvus_index_type_env = os.getenv("MILVUS_INDEX_TYPE")
if milvus_index_type_env:
milvus_index_type = milvus_index_type_env.upper()
log.info(f"Milvus index type from MILVUS_INDEX_TYPE env var: {milvus_index_type}")
else:
milvus_index_type = default_index_type
log.info(f"MILVUS_INDEX_TYPE env var not set, defaulting to: {milvus_index_type}")
index_creation_params = {} index_creation_params = {}
metric_type = os.getenv("MILVUS_METRIC_TYPE", "COSINE").upper() # Default to COSINE if index_type == "HNSW":
index_creation_params = {"M": MILVUS_HNSW_M, "efConstruction": MILVUS_HNSW_EFCONSTRUCTION}
if milvus_index_type == "HNSW": log.info(f"HNSW params: {index_creation_params}")
# Parameters for HNSW elif index_type == "IVF_FLAT":
m_env = os.getenv("MILVUS_HNSW_M", "16") index_creation_params = {"nlist": MILVUS_IVF_FLAT_NLIST}
ef_construction_env = os.getenv("MILVUS_HNSW_EFCONSTRUCTION", "100") log.info(f"IVF_FLAT params: {index_creation_params}")
try: elif index_type in ["FLAT", "AUTOINDEX"]:
m_val = int(m_env) log.info(f"Using {index_type} index with no specific build-time params.")
ef_val = int(ef_construction_env)
except ValueError:
log.warning(f"Invalid HNSW params M='{m_env}' or efConstruction='{ef_construction_env}'. Defaulting to M=16, efConstruction=100.")
m_val = 16
ef_val = 100
index_creation_params = {"M": m_val, "efConstruction": ef_val}
log.info(f"Using HNSW index with metric {metric_type}, params: {index_creation_params}")
elif milvus_index_type == "IVF_FLAT":
# Parameters for IVF_FLAT
nlist_env = os.getenv("MILVUS_IVF_FLAT_NLIST", "128")
try:
nlist = int(nlist_env)
except ValueError:
log.warning(f"Invalid MILVUS_IVF_FLAT_NLIST value '{nlist_env}'. Defaulting to 128.")
nlist = 128
index_creation_params = {"nlist": nlist}
log.info(f"Using IVF_FLAT index with metric {metric_type}, params: {index_creation_params}")
elif milvus_index_type == "FLAT":
log.info(f"Using FLAT index with metric {metric_type} (no specific build-time params).")
# No specific build-time parameters needed for FLAT
elif milvus_index_type == "AUTOINDEX":
log.info(f"Using AUTOINDEX with metric {metric_type} (params managed by Milvus).")
# No specific build-time parameters needed for AUTOINDEX
else: else:
log.warning( log.warning(
f"Unsupported or unrecognized MILVUS_INDEX_TYPE: '{milvus_index_type}'. " f"Unsupported MILVUS_INDEX_TYPE: '{index_type}'. "
f"Falling back to '{default_index_type}'. " f"Supported types: HNSW, IVF_FLAT, FLAT, AUTOINDEX. "
f"Supported types: HNSW, IVF_FLAT, FLAT, AUTOINDEX." f"Milvus will use its default for the collection if this type is not directly supported for index creation."
) )
milvus_index_type = default_index_type # Fallback to a safe default # For unsupported types, pass the type directly to Milvus; it might handle it or use a default.
# index_creation_params remains {} which is fine for AUTOINDEX/FLAT # If Milvus errors out, the user needs to correct the MILVUS_INDEX_TYPE env var.
log.info(f"Fell back to {default_index_type} index with metric {metric_type}.")
index_params.add_index( index_params.add_index(
field_name="vector", field_name="vector",
index_type=milvus_index_type, index_type=index_type,
metric_type=metric_type, metric_type=metric_type,
params=index_creation_params, params=index_creation_params,
) )
@ -175,7 +147,7 @@ class MilvusClient(VectorDBBase):
schema=schema, schema=schema,
index_params=index_params, index_params=index_params,
) )
log.info(f"Successfully created collection '{self.collection_prefix}_{collection_name}' with index type '{milvus_index_type}' and metric '{metric_type}'.") log.info(f"Successfully created collection '{self.collection_prefix}_{collection_name}' with index type '{index_type}' and metric '{metric_type}'.")
def has_collection(self, collection_name: str) -> bool: def has_collection(self, collection_name: str) -> bool: