68f663f2f4
主要变更: - 移除Hermes智能体及相关回调服务 - 新增知识库RAG、同步、调度、规范化和索引任务服务 - 重构orchestrator服务，增强运行时聊天功能 - 更新前端聊天、政策制度、设置等页面样式和逻辑 - 更新expense_claims和document_intelligence服务 - 删除llm_wiki相关服务和测试文件 - 更新docker-compose配置和启动脚本
1804 lines
70 KiB
Python
1804 lines
70 KiB
Python
import asyncio
|
|
import os
|
|
from typing import Any, final, Optional, Dict
|
|
from dataclasses import dataclass, fields
|
|
import numpy as np
|
|
from lightrag.utils import logger, compute_mdhash_id
|
|
from ..base import BaseVectorStorage
|
|
from ..constants import DEFAULT_MAX_FILE_PATH_LENGTH
|
|
from ..kg.shared_storage import get_data_init_lock
|
|
import pipmaster as pm
|
|
|
|
if not pm.is_installed("pymilvus"):
|
|
pm.install("pymilvus>=2.6.2")
|
|
|
|
import configparser
|
|
from pymilvus import MilvusClient, DataType, CollectionSchema, FieldSchema # type: ignore
|
|
from packaging import version
|
|
|
|
config = configparser.ConfigParser()
|
|
config.read("config.ini", "utf-8")
|
|
|
|
|
|
# Supported index types
|
|
SUPPORTED_INDEX_TYPES = {
|
|
"AUTOINDEX",
|
|
"HNSW",
|
|
"HNSW_SQ",
|
|
"HNSW_PQ",
|
|
"HNSW_PRQ",
|
|
"IVF_FLAT",
|
|
"IVF_SQ8",
|
|
"IVF_PQ",
|
|
"DISKANN",
|
|
"SCANN",
|
|
}
|
|
|
|
# Supported metric types
|
|
SUPPORTED_METRIC_TYPES = {"COSINE", "L2", "IP"}
|
|
|
|
# HNSW_SQ quantization types
|
|
SUPPORTED_SQ_TYPES = {"SQ4U", "SQ6", "SQ8", "BF16", "FP16"}
|
|
SUPPORTED_REFINE_TYPES = {"SQ6", "SQ8", "BF16", "FP16", "FP32"}
|
|
|
|
# Index type version requirements
|
|
# Important: HNSW_SQ was first introduced in Milvus 2.6.8 (not 2.5)
|
|
INDEX_VERSION_REQUIREMENTS = {
|
|
"HNSW_SQ": "2.6.8", # HNSW_SQ requires Milvus 2.6.8+ (supports sq_types such as SQ4U, SQ6, SQ8, BF16, FP16)
|
|
}
|
|
|
|
|
|
def _get_env_bool(key: str, default: bool = False) -> bool:
|
|
"""Parse environment variable as boolean"""
|
|
val = os.environ.get(key, "").lower()
|
|
if val in ("true", "1", "yes", "on"):
|
|
return True
|
|
elif val in ("false", "0", "no", "off"):
|
|
return False
|
|
return default
|
|
|
|
|
|
def _get_env_int(key: str, default: int) -> int:
|
|
"""Parse environment variable as integer"""
|
|
val = os.environ.get(key, "")
|
|
if val:
|
|
try:
|
|
return int(val)
|
|
except ValueError:
|
|
logger.warning(
|
|
f"Invalid integer value for {key}: {val}, using default {default}"
|
|
)
|
|
return default
|
|
|
|
|
|
@dataclass
|
|
class MilvusIndexConfig:
|
|
"""
|
|
Milvus vector index configuration class
|
|
|
|
Supports configuration via environment variables or initialization parameters.
|
|
Initialization parameters take precedence over environment variables.
|
|
"""
|
|
|
|
# Base configuration
|
|
index_type: Optional[str] = None
|
|
metric_type: Optional[str] = None
|
|
|
|
# HNSW series parameters
|
|
hnsw_m: Optional[int] = None
|
|
hnsw_ef_construction: Optional[int] = None
|
|
hnsw_ef: Optional[int] = None
|
|
|
|
# HNSW_SQ specific parameters
|
|
sq_type: Optional[str] = None
|
|
sq_refine: Optional[bool] = None
|
|
sq_refine_type: Optional[str] = None
|
|
sq_refine_k: Optional[int] = None
|
|
|
|
# IVF series parameters
|
|
ivf_nlist: Optional[int] = None
|
|
ivf_nprobe: Optional[int] = None
|
|
|
|
def __post_init__(self):
|
|
"""Load configuration from environment variables (init parameters take precedence)"""
|
|
# Index type
|
|
self.index_type = (
|
|
self.index_type or os.environ.get("MILVUS_INDEX_TYPE", "AUTOINDEX")
|
|
).upper()
|
|
|
|
# Metric type
|
|
self.metric_type = (
|
|
self.metric_type or os.environ.get("MILVUS_METRIC_TYPE", "COSINE")
|
|
).upper()
|
|
|
|
# HNSW parameters
|
|
# Defaults aligned with Milvus 2.4+ official documentation
|
|
if self.hnsw_m is None:
|
|
self.hnsw_m = _get_env_int("MILVUS_HNSW_M", 16)
|
|
if self.hnsw_ef_construction is None:
|
|
self.hnsw_ef_construction = _get_env_int("MILVUS_HNSW_EF_CONSTRUCTION", 360)
|
|
if self.hnsw_ef is None:
|
|
self.hnsw_ef = _get_env_int("MILVUS_HNSW_EF", 200)
|
|
|
|
# HNSW_SQ parameters
|
|
if self.sq_type is None:
|
|
self.sq_type = os.environ.get("MILVUS_HNSW_SQ_TYPE", "SQ8").upper()
|
|
if self.sq_refine is None:
|
|
self.sq_refine = _get_env_bool("MILVUS_HNSW_SQ_REFINE", False)
|
|
if self.sq_refine_type is None:
|
|
self.sq_refine_type = os.environ.get(
|
|
"MILVUS_HNSW_SQ_REFINE_TYPE", "FP32"
|
|
).upper()
|
|
if self.sq_refine_k is None:
|
|
self.sq_refine_k = _get_env_int("MILVUS_HNSW_SQ_REFINE_K", 10)
|
|
|
|
# IVF parameters
|
|
if self.ivf_nlist is None:
|
|
self.ivf_nlist = _get_env_int("MILVUS_IVF_NLIST", 1024)
|
|
if self.ivf_nprobe is None:
|
|
self.ivf_nprobe = _get_env_int("MILVUS_IVF_NPROBE", 16)
|
|
|
|
# Validate configuration
|
|
self._validate()
|
|
|
|
def _validate(self):
|
|
"""Validate configuration validity"""
|
|
if self.index_type not in SUPPORTED_INDEX_TYPES:
|
|
raise ValueError(
|
|
f"Unsupported index type: {self.index_type}. "
|
|
f"Supported: {SUPPORTED_INDEX_TYPES}"
|
|
)
|
|
|
|
if self.metric_type not in SUPPORTED_METRIC_TYPES:
|
|
raise ValueError(
|
|
f"Unsupported metric type: {self.metric_type}. "
|
|
f"Supported: {SUPPORTED_METRIC_TYPES}"
|
|
)
|
|
|
|
if self.index_type == "HNSW_SQ":
|
|
if self.sq_type not in SUPPORTED_SQ_TYPES:
|
|
raise ValueError(
|
|
f"Unsupported sq_type: {self.sq_type}. "
|
|
f"Supported: {SUPPORTED_SQ_TYPES}"
|
|
)
|
|
if self.sq_refine and self.sq_refine_type not in SUPPORTED_REFINE_TYPES:
|
|
raise ValueError(
|
|
f"Unsupported refine_type: {self.sq_refine_type}. "
|
|
f"Supported: {SUPPORTED_REFINE_TYPES}"
|
|
)
|
|
|
|
# Parameter range validation
|
|
if not (2 <= self.hnsw_m <= 2048):
|
|
raise ValueError(f"hnsw_m must be in [2, 2048], got {self.hnsw_m}")
|
|
if self.hnsw_ef_construction < 1:
|
|
raise ValueError(
|
|
f"hnsw_ef_construction must be >= 1, got {self.hnsw_ef_construction}"
|
|
)
|
|
if self.ivf_nlist < 1 or self.ivf_nlist > 65536:
|
|
raise ValueError(f"ivf_nlist must be in [1, 65536], got {self.ivf_nlist}")
|
|
|
|
def validate_milvus_version(self, server_version: str) -> None:
|
|
"""
|
|
Validate Milvus server version supports the configured index type
|
|
|
|
Args:
|
|
server_version: Milvus server version string (e.g., "2.6.9")
|
|
|
|
Raises:
|
|
ValueError: Version does not meet index type requirements
|
|
"""
|
|
current_ver = version.parse(
|
|
server_version.split("-")[0]
|
|
) # Handle "2.6.9-dev" format
|
|
|
|
# Check HNSW_SQ index type version requirements (requires 2.6.8+)
|
|
if self.index_type == "HNSW_SQ":
|
|
required = INDEX_VERSION_REQUIREMENTS["HNSW_SQ"]
|
|
if current_ver < version.parse(required):
|
|
raise ValueError(
|
|
f"HNSW_SQ requires Milvus {required}+, "
|
|
f"current version: {server_version}"
|
|
)
|
|
|
|
logger.info(
|
|
f"Milvus version {server_version} validated for index type "
|
|
f"{self.index_type}"
|
|
+ (f" with sq_type {self.sq_type}" if self.index_type == "HNSW_SQ" else "")
|
|
)
|
|
|
|
def build_index_params(self, index_params, field_name: str = "vector"):
|
|
"""
|
|
Build pymilvus index parameters
|
|
|
|
Args:
|
|
index_params: IndexParams instance (from compatibility helper or client.prepare_index_params())
|
|
field_name: Vector field name
|
|
|
|
Returns:
|
|
IndexParams object, or a dict fallback when direct API creation is needed.
|
|
"""
|
|
if index_params is None:
|
|
if self.index_type == "AUTOINDEX":
|
|
logger.info(
|
|
"Using AUTOINDEX with direct API fallback because IndexParams is unavailable"
|
|
)
|
|
return {
|
|
"field_name": field_name,
|
|
"index_type": self.index_type,
|
|
"metric_type": self.metric_type,
|
|
"params": {},
|
|
}
|
|
raise RuntimeError(
|
|
f"IndexParams not available but required for index type "
|
|
f"'{self.index_type}'. Ensure pymilvus is installed correctly."
|
|
)
|
|
|
|
params: Dict[str, Any] = {}
|
|
|
|
# HNSW series indexes
|
|
if self.index_type in ("HNSW", "HNSW_SQ", "HNSW_PQ", "HNSW_PRQ"):
|
|
params["M"] = self.hnsw_m
|
|
params["efConstruction"] = self.hnsw_ef_construction
|
|
|
|
# HNSW_SQ specific parameters
|
|
if self.index_type == "HNSW_SQ":
|
|
params["sq_type"] = self.sq_type
|
|
if self.sq_refine:
|
|
params["refine"] = True
|
|
params["refine_type"] = self.sq_refine_type
|
|
|
|
# IVF series indexes
|
|
elif self.index_type in ("IVF_FLAT", "IVF_SQ8", "IVF_PQ"):
|
|
params["nlist"] = self.ivf_nlist
|
|
|
|
# DISKANN / SCANN have no additional params
|
|
|
|
index_params.add_index(
|
|
field_name=field_name,
|
|
index_type=self.index_type,
|
|
metric_type=self.metric_type,
|
|
params=params,
|
|
)
|
|
|
|
logger.info(
|
|
f"Milvus index configured: type={self.index_type}, "
|
|
f"metric={self.metric_type}, params={params}"
|
|
)
|
|
|
|
return index_params
|
|
|
|
def build_search_params(self) -> Dict[str, Any]:
|
|
"""
|
|
Build search parameters
|
|
|
|
Returns:
|
|
Search parameters dictionary
|
|
"""
|
|
search_params: Dict[str, Any] = {}
|
|
|
|
if self.index_type in ("HNSW", "HNSW_SQ", "HNSW_PQ", "HNSW_PRQ"):
|
|
search_params["ef"] = self.hnsw_ef
|
|
if self.index_type == "HNSW_SQ" and self.sq_refine:
|
|
search_params["refine_k"] = self.sq_refine_k
|
|
|
|
elif self.index_type in ("IVF_FLAT", "IVF_SQ8", "IVF_PQ"):
|
|
search_params["nprobe"] = self.ivf_nprobe
|
|
|
|
return {"params": search_params} if search_params else {}
|
|
|
|
@classmethod
|
|
def get_config_field_names(cls) -> set:
|
|
"""Get all configuration field names from the dataclass.
|
|
|
|
This method provides a single source of truth for configuration parameter names,
|
|
eliminating the need to maintain duplicate hardcoded lists elsewhere.
|
|
|
|
Returns:
|
|
Set of field names that can be used to extract configuration from kwargs
|
|
"""
|
|
return {f.name for f in fields(cls)}
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Export configuration as dictionary (for logging/debugging)"""
|
|
return {
|
|
"index_type": self.index_type,
|
|
"metric_type": self.metric_type,
|
|
"hnsw_m": self.hnsw_m,
|
|
"hnsw_ef_construction": self.hnsw_ef_construction,
|
|
"hnsw_ef": self.hnsw_ef,
|
|
"sq_type": self.sq_type if self.index_type == "HNSW_SQ" else None,
|
|
"sq_refine": self.sq_refine if self.index_type == "HNSW_SQ" else None,
|
|
"sq_refine_type": (
|
|
self.sq_refine_type
|
|
if self.index_type == "HNSW_SQ" and self.sq_refine
|
|
else None
|
|
),
|
|
"sq_refine_k": (
|
|
self.sq_refine_k
|
|
if self.index_type == "HNSW_SQ" and self.sq_refine
|
|
else None
|
|
),
|
|
"ivf_nlist": (
|
|
self.ivf_nlist if self.index_type.startswith("IVF") else None
|
|
),
|
|
"ivf_nprobe": (
|
|
self.ivf_nprobe if self.index_type.startswith("IVF") else None
|
|
),
|
|
}
|
|
|
|
|
|
@final
|
|
@dataclass
|
|
class MilvusVectorDBStorage(BaseVectorStorage):
|
|
def _get_milvus_connection_kwargs(self, include_db_name: bool = True) -> dict:
|
|
"""Build Milvus connection kwargs from env/config."""
|
|
connection_kwargs = {
|
|
"uri": os.environ.get(
|
|
"MILVUS_URI",
|
|
config.get(
|
|
"milvus",
|
|
"uri",
|
|
fallback=os.path.join(
|
|
self.global_config["working_dir"], "milvus_lite.db"
|
|
),
|
|
),
|
|
),
|
|
"user": os.environ.get(
|
|
"MILVUS_USER", config.get("milvus", "user", fallback=None)
|
|
),
|
|
"password": os.environ.get(
|
|
"MILVUS_PASSWORD",
|
|
config.get("milvus", "password", fallback=None),
|
|
),
|
|
"token": os.environ.get(
|
|
"MILVUS_TOKEN", config.get("milvus", "token", fallback=None)
|
|
),
|
|
}
|
|
|
|
db_name = os.environ.get(
|
|
"MILVUS_DB_NAME",
|
|
config.get("milvus", "db_name", fallback=None),
|
|
)
|
|
if include_db_name and db_name:
|
|
connection_kwargs["db_name"] = db_name
|
|
|
|
return connection_kwargs
|
|
|
|
def _get_milvus_db_name(self) -> Optional[str]:
|
|
"""Return the configured Milvus database name, if any."""
|
|
db_name = self._get_milvus_connection_kwargs(include_db_name=True).get(
|
|
"db_name"
|
|
)
|
|
if db_name is None:
|
|
return None
|
|
|
|
normalized_name = str(db_name).strip()
|
|
return normalized_name or None
|
|
|
|
def _create_milvus_client(self) -> MilvusClient:
|
|
"""Create a Milvus client and ensure the configured database exists."""
|
|
client = MilvusClient(
|
|
**self._get_milvus_connection_kwargs(include_db_name=False)
|
|
)
|
|
db_name = self._get_milvus_db_name()
|
|
|
|
if not db_name:
|
|
return client
|
|
|
|
existing_databases = set(client.list_databases())
|
|
if db_name not in existing_databases:
|
|
logger.warning(
|
|
f"[{self.workspace}] Milvus database '{db_name}' not found, creating it"
|
|
)
|
|
client.create_database(db_name)
|
|
|
|
use_database = getattr(client, "use_database", None) or getattr(
|
|
client, "using_database", None
|
|
)
|
|
if callable(use_database):
|
|
use_database(db_name)
|
|
logger.debug(
|
|
f"[{self.workspace}] Using Milvus database '{db_name}' for namespace '{self.namespace}'"
|
|
)
|
|
return client
|
|
|
|
return MilvusClient(**self._get_milvus_connection_kwargs(include_db_name=True))
|
|
|
|
def _create_schema_for_namespace(self) -> CollectionSchema:
|
|
"""Create schema based on the current instance's namespace"""
|
|
|
|
# Get vector dimension from embedding_func
|
|
dimension = self.embedding_func.embedding_dim
|
|
|
|
# Base fields (common to all collections)
|
|
base_fields = [
|
|
FieldSchema(
|
|
name="id", dtype=DataType.VARCHAR, max_length=64, is_primary=True
|
|
),
|
|
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension),
|
|
FieldSchema(name="created_at", dtype=DataType.INT64),
|
|
]
|
|
|
|
# Determine specific fields based on namespace
|
|
if self.namespace.endswith("entities"):
|
|
specific_fields = [
|
|
FieldSchema(
|
|
name="entity_name",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=512,
|
|
nullable=True,
|
|
),
|
|
FieldSchema(
|
|
name="file_path",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=DEFAULT_MAX_FILE_PATH_LENGTH,
|
|
nullable=True,
|
|
),
|
|
]
|
|
description = "LightRAG entities vector storage"
|
|
|
|
elif self.namespace.endswith("relationships"):
|
|
specific_fields = [
|
|
FieldSchema(
|
|
name="src_id", dtype=DataType.VARCHAR, max_length=512, nullable=True
|
|
),
|
|
FieldSchema(
|
|
name="tgt_id", dtype=DataType.VARCHAR, max_length=512, nullable=True
|
|
),
|
|
FieldSchema(
|
|
name="file_path",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=DEFAULT_MAX_FILE_PATH_LENGTH,
|
|
nullable=True,
|
|
),
|
|
]
|
|
description = "LightRAG relationships vector storage"
|
|
|
|
elif self.namespace.endswith("chunks"):
|
|
specific_fields = [
|
|
FieldSchema(
|
|
name="full_doc_id",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=64,
|
|
nullable=True,
|
|
),
|
|
FieldSchema(
|
|
name="file_path",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=DEFAULT_MAX_FILE_PATH_LENGTH,
|
|
nullable=True,
|
|
),
|
|
]
|
|
description = "LightRAG chunks vector storage"
|
|
|
|
else:
|
|
# Default generic schema (backward compatibility)
|
|
specific_fields = [
|
|
FieldSchema(
|
|
name="file_path",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=DEFAULT_MAX_FILE_PATH_LENGTH,
|
|
nullable=True,
|
|
),
|
|
]
|
|
description = "LightRAG generic vector storage"
|
|
|
|
# Merge all fields
|
|
all_fields = base_fields + specific_fields
|
|
|
|
return CollectionSchema(
|
|
fields=all_fields,
|
|
description=description,
|
|
enable_dynamic_field=True, # Support dynamic fields
|
|
)
|
|
|
|
def _get_index_params(self):
|
|
"""Get IndexParams in a version-compatible way"""
|
|
try:
|
|
# Try to use client's prepare_index_params method (most common)
|
|
if hasattr(self._client, "prepare_index_params"):
|
|
return self._client.prepare_index_params()
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
# Try to import IndexParams from different possible locations
|
|
from pymilvus.client.prepare import IndexParams # type: ignore
|
|
|
|
return IndexParams()
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
from pymilvus.client.types import IndexParams # type: ignore
|
|
|
|
return IndexParams()
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
from pymilvus import IndexParams # type: ignore
|
|
|
|
return IndexParams()
|
|
except ImportError:
|
|
pass
|
|
|
|
# If all else fails, return None to use fallback method
|
|
return None
|
|
|
|
def _create_scalar_index_fallback(self, field_name: str, index_type: str):
|
|
"""Fallback method to create scalar index using direct API"""
|
|
# Skip unsupported index types
|
|
if index_type == "SORTED":
|
|
logger.info(
|
|
f"[{self.workspace}] Skipping SORTED index for {field_name} (not supported in this Milvus version)"
|
|
)
|
|
return
|
|
|
|
try:
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
field_name=field_name,
|
|
index_params={"index_type": index_type},
|
|
)
|
|
logger.debug(
|
|
f"[{self.workspace}] Created {field_name} index using fallback method"
|
|
)
|
|
except Exception as e:
|
|
logger.info(
|
|
f"[{self.workspace}] Could not create {field_name} index using fallback method: {e}"
|
|
)
|
|
|
|
def _create_indexes_after_collection(self):
|
|
"""Create indexes after collection is created"""
|
|
# Build vector index using index configuration
|
|
# Use compatibility helper to get IndexParams
|
|
index_params_for_vector = self._get_index_params()
|
|
|
|
vector_index_params = self.index_config.build_index_params(
|
|
index_params_for_vector, field_name="vector"
|
|
)
|
|
|
|
# Re-raise exceptions to surface vector index creation failures
|
|
if isinstance(vector_index_params, dict):
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
field_name=vector_index_params["field_name"],
|
|
index_params={
|
|
"index_type": vector_index_params["index_type"],
|
|
"metric_type": vector_index_params["metric_type"],
|
|
"params": vector_index_params["params"],
|
|
},
|
|
)
|
|
else:
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
index_params=vector_index_params,
|
|
)
|
|
|
|
logger.debug(
|
|
f"[{self.workspace}] Created vector index with config: {self.index_config.to_dict()}"
|
|
)
|
|
|
|
# Create scalar indexes based on namespace
|
|
# Wrap scalar index creation in try-except to allow graceful degradation
|
|
try:
|
|
# Try to get IndexParams in a version-compatible way
|
|
scalar_index_params = self._get_index_params()
|
|
|
|
if scalar_index_params is not None:
|
|
# Create scalar indexes based on namespace
|
|
if self.namespace.endswith("entities"):
|
|
# Create indexes for entity fields
|
|
try:
|
|
entity_name_index = self._get_index_params()
|
|
entity_name_index.add_index(
|
|
field_name="entity_name", index_type="INVERTED"
|
|
)
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
index_params=entity_name_index,
|
|
)
|
|
except Exception as e:
|
|
logger.debug(
|
|
f"[{self.workspace}] IndexParams method failed for entity_name: {e}"
|
|
)
|
|
self._create_scalar_index_fallback("entity_name", "INVERTED")
|
|
|
|
elif self.namespace.endswith("relationships"):
|
|
# Create indexes for relationship fields
|
|
try:
|
|
src_id_index = self._get_index_params()
|
|
src_id_index.add_index(
|
|
field_name="src_id", index_type="INVERTED"
|
|
)
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
index_params=src_id_index,
|
|
)
|
|
except Exception as e:
|
|
logger.debug(
|
|
f"[{self.workspace}] IndexParams method failed for src_id: {e}"
|
|
)
|
|
self._create_scalar_index_fallback("src_id", "INVERTED")
|
|
|
|
try:
|
|
tgt_id_index = self._get_index_params()
|
|
tgt_id_index.add_index(
|
|
field_name="tgt_id", index_type="INVERTED"
|
|
)
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
index_params=tgt_id_index,
|
|
)
|
|
except Exception as e:
|
|
logger.debug(
|
|
f"[{self.workspace}] IndexParams method failed for tgt_id: {e}"
|
|
)
|
|
self._create_scalar_index_fallback("tgt_id", "INVERTED")
|
|
|
|
elif self.namespace.endswith("chunks"):
|
|
# Create indexes for chunk fields
|
|
try:
|
|
doc_id_index = self._get_index_params()
|
|
doc_id_index.add_index(
|
|
field_name="full_doc_id", index_type="INVERTED"
|
|
)
|
|
self._client.create_index(
|
|
collection_name=self.final_namespace,
|
|
index_params=doc_id_index,
|
|
)
|
|
except Exception as e:
|
|
logger.debug(
|
|
f"[{self.workspace}] IndexParams method failed for full_doc_id: {e}"
|
|
)
|
|
self._create_scalar_index_fallback("full_doc_id", "INVERTED")
|
|
|
|
else:
|
|
# Fallback to direct API calls if IndexParams is not available
|
|
logger.info(
|
|
f"[{self.workspace}] IndexParams not available, using fallback methods for {self.namespace}"
|
|
)
|
|
|
|
# Create scalar indexes using fallback
|
|
if self.namespace.endswith("entities"):
|
|
self._create_scalar_index_fallback("entity_name", "INVERTED")
|
|
elif self.namespace.endswith("relationships"):
|
|
self._create_scalar_index_fallback("src_id", "INVERTED")
|
|
self._create_scalar_index_fallback("tgt_id", "INVERTED")
|
|
elif self.namespace.endswith("chunks"):
|
|
self._create_scalar_index_fallback("full_doc_id", "INVERTED")
|
|
|
|
logger.info(
|
|
f"[{self.workspace}] Created indexes for collection: {self.namespace}"
|
|
)
|
|
|
|
except Exception as e:
|
|
# Scalar index failures are logged as warnings (not critical)
|
|
logger.warning(
|
|
f"[{self.workspace}] Failed to create some scalar indexes for {self.namespace}: {e}"
|
|
)
|
|
|
|
def _get_required_fields_for_namespace(self) -> dict:
|
|
"""Get required core field definitions for current namespace"""
|
|
|
|
# Base fields (common to all types)
|
|
base_fields = {
|
|
"id": {"type": "VarChar", "is_primary": True},
|
|
"vector": {"type": "FloatVector"},
|
|
"created_at": {"type": "Int64"},
|
|
}
|
|
|
|
# Add specific fields based on namespace
|
|
if self.namespace.endswith("entities"):
|
|
specific_fields = {
|
|
"entity_name": {"type": "VarChar"},
|
|
"file_path": {"type": "VarChar"},
|
|
}
|
|
elif self.namespace.endswith("relationships"):
|
|
specific_fields = {
|
|
"src_id": {"type": "VarChar"},
|
|
"tgt_id": {"type": "VarChar"},
|
|
"file_path": {"type": "VarChar"},
|
|
}
|
|
elif self.namespace.endswith("chunks"):
|
|
specific_fields = {
|
|
"full_doc_id": {"type": "VarChar"},
|
|
"file_path": {"type": "VarChar"},
|
|
}
|
|
else:
|
|
specific_fields = {
|
|
"file_path": {"type": "VarChar"},
|
|
}
|
|
|
|
return {**base_fields, **specific_fields}
|
|
|
|
def _is_field_compatible(self, existing_field: dict, expected_config: dict) -> bool:
|
|
"""Check compatibility of a single field"""
|
|
field_name = existing_field.get("name", "unknown")
|
|
existing_type = existing_field.get("type")
|
|
expected_type = expected_config.get("type")
|
|
|
|
logger.debug(
|
|
f"[{self.workspace}] Checking field '{field_name}': existing_type={existing_type} (type={type(existing_type)}), expected_type={expected_type}"
|
|
)
|
|
|
|
# Convert DataType enum values to string names if needed
|
|
original_existing_type = existing_type
|
|
if hasattr(existing_type, "name"):
|
|
existing_type = existing_type.name
|
|
logger.debug(
|
|
f"[{self.workspace}] Converted enum to name: {original_existing_type} -> {existing_type}"
|
|
)
|
|
elif isinstance(existing_type, int):
|
|
# Map common Milvus internal type codes to type names for backward compatibility
|
|
type_mapping = {
|
|
21: "VarChar",
|
|
101: "FloatVector",
|
|
5: "Int64",
|
|
9: "Double",
|
|
}
|
|
mapped_type = type_mapping.get(existing_type, str(existing_type))
|
|
logger.debug(
|
|
f"[{self.workspace}] Mapped numeric type: {existing_type} -> {mapped_type}"
|
|
)
|
|
existing_type = mapped_type
|
|
|
|
# Normalize type names for comparison
|
|
type_aliases = {
|
|
"VARCHAR": "VarChar",
|
|
"String": "VarChar",
|
|
"FLOAT_VECTOR": "FloatVector",
|
|
"INT64": "Int64",
|
|
"BigInt": "Int64",
|
|
"DOUBLE": "Double",
|
|
"Float": "Double",
|
|
}
|
|
|
|
original_existing = existing_type
|
|
original_expected = expected_type
|
|
existing_type = type_aliases.get(existing_type, existing_type)
|
|
expected_type = type_aliases.get(expected_type, expected_type)
|
|
|
|
if original_existing != existing_type or original_expected != expected_type:
|
|
logger.debug(
|
|
f"[{self.workspace}] Applied aliases: {original_existing} -> {existing_type}, {original_expected} -> {expected_type}"
|
|
)
|
|
|
|
# Basic type compatibility check
|
|
type_compatible = existing_type == expected_type
|
|
logger.debug(
|
|
f"[{self.workspace}] Type compatibility for '{field_name}': {existing_type} == {expected_type} -> {type_compatible}"
|
|
)
|
|
|
|
if not type_compatible:
|
|
logger.warning(
|
|
f"[{self.workspace}] Type mismatch for field '{field_name}': expected {expected_type}, got {existing_type}"
|
|
)
|
|
return False
|
|
|
|
# Primary key check - be more flexible about primary key detection
|
|
if expected_config.get("is_primary"):
|
|
# Check multiple possible field names for primary key status
|
|
is_primary = (
|
|
existing_field.get("is_primary_key", False)
|
|
or existing_field.get("is_primary", False)
|
|
or existing_field.get("primary_key", False)
|
|
)
|
|
logger.debug(
|
|
f"[{self.workspace}] Primary key check for '{field_name}': expected=True, actual={is_primary}"
|
|
)
|
|
logger.debug(
|
|
f"[{self.workspace}] Raw field data for '{field_name}': {existing_field}"
|
|
)
|
|
|
|
# For ID field, be more lenient - if it's the ID field, assume it should be primary
|
|
if field_name == "id" and not is_primary:
|
|
logger.info(
|
|
f"[{self.workspace}] ID field '{field_name}' not marked as primary in existing collection, but treating as compatible"
|
|
)
|
|
# Don't fail for ID field primary key mismatch
|
|
elif not is_primary:
|
|
logger.warning(
|
|
f"[{self.workspace}] Primary key mismatch for field '{field_name}': expected primary key, but field is not primary"
|
|
)
|
|
return False
|
|
|
|
logger.debug(f"[{self.workspace}] Field '{field_name}' is compatible")
|
|
return True
|
|
|
|
def _check_vector_dimension(self, collection_info: dict):
|
|
"""Check vector dimension compatibility"""
|
|
current_dimension = self.embedding_func.embedding_dim
|
|
|
|
# Find vector field dimension
|
|
for field in collection_info.get("fields", []):
|
|
if field.get("name") == "vector":
|
|
field_type = field.get("type")
|
|
|
|
# Extract type name from DataType enum or string
|
|
type_name = None
|
|
if hasattr(field_type, "name"):
|
|
type_name = field_type.name
|
|
elif isinstance(field_type, str):
|
|
type_name = field_type
|
|
else:
|
|
type_name = str(field_type)
|
|
|
|
# Check if it's a vector type (supports multiple formats)
|
|
if type_name in ["FloatVector", "FLOAT_VECTOR"]:
|
|
existing_dimension = field.get("params", {}).get("dim")
|
|
|
|
# Convert both to int for comparison to handle type mismatches
|
|
# (Milvus API may return string "1024" vs int 1024)
|
|
try:
|
|
existing_dim_int = (
|
|
int(existing_dimension)
|
|
if existing_dimension is not None
|
|
else None
|
|
)
|
|
current_dim_int = (
|
|
int(current_dimension)
|
|
if current_dimension is not None
|
|
else None
|
|
)
|
|
except (TypeError, ValueError) as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Failed to parse dimensions: existing={existing_dimension} (type={type(existing_dimension)}), "
|
|
f"current={current_dimension} (type={type(current_dimension)}), error={e}"
|
|
)
|
|
raise ValueError(
|
|
f"Invalid dimension values for collection '{self.final_namespace}': "
|
|
f"existing={existing_dimension}, current={current_dimension}"
|
|
) from e
|
|
|
|
if existing_dim_int != current_dim_int:
|
|
raise ValueError(
|
|
f"Vector dimension mismatch for collection '{self.final_namespace}': "
|
|
f"existing={existing_dim_int}, current={current_dim_int}"
|
|
)
|
|
|
|
logger.debug(
|
|
f"[{self.workspace}] Vector dimension check passed: {current_dim_int}"
|
|
)
|
|
return
|
|
|
|
# If no vector field found, this might be an old collection created with simple schema
|
|
logger.warning(
|
|
f"[{self.workspace}] Vector field not found in collection '{self.namespace}'. This might be an old collection created with simple schema."
|
|
)
|
|
logger.warning(
|
|
f"[{self.workspace}] Consider recreating the collection for optimal performance."
|
|
)
|
|
return
|
|
|
|
def _check_file_path_length_restriction(self, collection_info: dict) -> bool:
|
|
"""Check if collection has file_path length restrictions that need migration
|
|
|
|
Returns:
|
|
bool: True if migration is needed, False otherwise
|
|
"""
|
|
existing_fields = {
|
|
field["name"]: field for field in collection_info.get("fields", [])
|
|
}
|
|
|
|
# Check if file_path field exists and has length restrictions
|
|
if "file_path" in existing_fields:
|
|
file_path_field = existing_fields["file_path"]
|
|
# Get max_length from field params
|
|
max_length = file_path_field.get("params", {}).get("max_length")
|
|
|
|
if max_length and max_length < DEFAULT_MAX_FILE_PATH_LENGTH:
|
|
logger.info(
|
|
f"[{self.workspace}] Collection {self.namespace} has file_path max_length={max_length}, "
|
|
f"needs migration to {DEFAULT_MAX_FILE_PATH_LENGTH}"
|
|
)
|
|
return True
|
|
|
|
return False
|
|
|
|
def _check_schema_compatibility(self, collection_info: dict):
|
|
"""Check schema field compatibility and detect migration needs"""
|
|
existing_fields = {
|
|
field["name"]: field for field in collection_info.get("fields", [])
|
|
}
|
|
|
|
# Check if this is an old collection created with simple schema
|
|
has_vector_field = any(
|
|
field.get("name") == "vector" for field in collection_info.get("fields", [])
|
|
)
|
|
|
|
if not has_vector_field:
|
|
logger.warning(
|
|
f"[{self.workspace}] Collection {self.namespace} appears to be created with old simple schema (no vector field)"
|
|
)
|
|
logger.warning(
|
|
f"[{self.workspace}] This collection will work but may have suboptimal performance"
|
|
)
|
|
logger.warning(
|
|
f"[{self.workspace}] Consider recreating the collection for optimal performance"
|
|
)
|
|
return
|
|
|
|
# Check if migration is needed for file_path length restrictions
|
|
if self._check_file_path_length_restriction(collection_info):
|
|
logger.info(
|
|
f"[{self.workspace}] Starting automatic migration for collection {self.namespace}"
|
|
)
|
|
self._migrate_collection_schema()
|
|
return
|
|
|
|
# For collections with vector field, check basic compatibility
|
|
# Only check for critical incompatibilities, not missing optional fields
|
|
critical_fields = {"id": {"type": "VarChar", "is_primary": True}}
|
|
|
|
incompatible_fields = []
|
|
|
|
for field_name, expected_config in critical_fields.items():
|
|
if field_name in existing_fields:
|
|
existing_field = existing_fields[field_name]
|
|
if not self._is_field_compatible(existing_field, expected_config):
|
|
incompatible_fields.append(
|
|
f"{field_name}: expected {expected_config['type']}, "
|
|
f"got {existing_field.get('type')}"
|
|
)
|
|
|
|
if incompatible_fields:
|
|
raise ValueError(
|
|
f"Critical schema incompatibility in collection '{self.final_namespace}': {incompatible_fields}"
|
|
)
|
|
|
|
# Get all expected fields for informational purposes
|
|
expected_fields = self._get_required_fields_for_namespace()
|
|
missing_fields = [
|
|
field for field in expected_fields if field not in existing_fields
|
|
]
|
|
|
|
if missing_fields:
|
|
logger.info(
|
|
f"[{self.workspace}] Collection {self.namespace} missing optional fields: {missing_fields}"
|
|
)
|
|
logger.info(
|
|
"These fields would be available in a newly created collection for better performance"
|
|
)
|
|
|
|
logger.debug(
|
|
f"[{self.workspace}] Schema compatibility check passed for {self.namespace}"
|
|
)
|
|
|
|
def _migrate_collection_schema(self):
|
|
"""Migrate collection schema using query_iterator - completely solves query window limitations"""
|
|
original_collection_name = self.final_namespace
|
|
temp_collection_name = f"{self.final_namespace}_temp"
|
|
iterator = None
|
|
|
|
try:
|
|
logger.info(
|
|
f"[{self.workspace}] Starting iterator-based schema migration for {self.namespace}"
|
|
)
|
|
|
|
# Step 1: Create temporary collection with new schema
|
|
logger.info(
|
|
f"[{self.workspace}] Step 1: Creating temporary collection: {temp_collection_name}"
|
|
)
|
|
# Temporarily update final_namespace for index creation
|
|
self.final_namespace = temp_collection_name
|
|
new_schema = self._create_schema_for_namespace()
|
|
self._client.create_collection(
|
|
collection_name=temp_collection_name, schema=new_schema
|
|
)
|
|
try:
|
|
self._create_indexes_after_collection()
|
|
except Exception as index_error:
|
|
logger.warning(
|
|
f"[{self.workspace}] Failed to create indexes for new collection: {index_error}"
|
|
)
|
|
# Continue with migration even if index creation fails
|
|
|
|
# Load the new collection
|
|
self._client.load_collection(temp_collection_name)
|
|
|
|
# Step 2: Copy data using query_iterator (solves query window limitation)
|
|
logger.info(
|
|
f"[{self.workspace}] Step 2: Copying data using query_iterator from: {original_collection_name}"
|
|
)
|
|
|
|
# Create query iterator
|
|
try:
|
|
iterator = self._client.query_iterator(
|
|
collection_name=original_collection_name,
|
|
batch_size=2000, # Adjustable batch size for optimal performance
|
|
output_fields=["*"], # Get all fields
|
|
)
|
|
logger.debug(f"[{self.workspace}] Query iterator created successfully")
|
|
except Exception as iterator_error:
|
|
logger.error(
|
|
f"[{self.workspace}] Failed to create query iterator: {iterator_error}"
|
|
)
|
|
raise
|
|
|
|
# Iterate through all data
|
|
total_migrated = 0
|
|
batch_number = 1
|
|
|
|
while True:
|
|
try:
|
|
batch_data = iterator.next()
|
|
if not batch_data:
|
|
# No more data available
|
|
break
|
|
|
|
# Insert batch data to new collection
|
|
try:
|
|
self._client.insert(
|
|
collection_name=temp_collection_name, data=batch_data
|
|
)
|
|
total_migrated += len(batch_data)
|
|
|
|
logger.info(
|
|
f"[{self.workspace}] Iterator batch {batch_number}: "
|
|
f"processed {len(batch_data)} records, total migrated: {total_migrated}"
|
|
)
|
|
batch_number += 1
|
|
|
|
except Exception as batch_error:
|
|
logger.error(
|
|
f"[{self.workspace}] Failed to insert iterator batch {batch_number}: {batch_error}"
|
|
)
|
|
raise
|
|
|
|
except Exception as next_error:
|
|
logger.error(
|
|
f"[{self.workspace}] Iterator next() failed at batch {batch_number}: {next_error}"
|
|
)
|
|
raise
|
|
|
|
if total_migrated > 0:
|
|
logger.info(
|
|
f"[{self.workspace}] Successfully migrated {total_migrated} records using iterator"
|
|
)
|
|
else:
|
|
logger.info(
|
|
f"[{self.workspace}] No data found in original collection, migration completed"
|
|
)
|
|
|
|
# Step 3: Rename origin collection (keep for safety)
|
|
logger.info(
|
|
f"[{self.workspace}] Step 3: Rename origin collection to {original_collection_name}_old"
|
|
)
|
|
try:
|
|
self._client.rename_collection(
|
|
original_collection_name, f"{original_collection_name}_old"
|
|
)
|
|
except Exception as rename_error:
|
|
try:
|
|
logger.warning(
|
|
f"[{self.workspace}] Try to drop origin collection instead"
|
|
)
|
|
self._client.drop_collection(original_collection_name)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Rename operation failed: {rename_error}"
|
|
)
|
|
raise e
|
|
|
|
# Step 4: Rename temporary collection to original name
|
|
logger.info(
|
|
f"[{self.workspace}] Step 4: Renaming collection {temp_collection_name} -> {original_collection_name}"
|
|
)
|
|
try:
|
|
self._client.rename_collection(
|
|
temp_collection_name, original_collection_name
|
|
)
|
|
logger.info(f"[{self.workspace}] Rename operation completed")
|
|
except Exception as rename_error:
|
|
logger.error(
|
|
f"[{self.workspace}] Rename operation failed: {rename_error}"
|
|
)
|
|
raise RuntimeError(
|
|
f"Failed to rename collection: {rename_error}"
|
|
) from rename_error
|
|
|
|
# Restore final_namespace
|
|
self.final_namespace = original_collection_name
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Iterator-based migration failed for {self.namespace}: {e}"
|
|
)
|
|
|
|
# Attempt cleanup of temporary collection if it exists
|
|
try:
|
|
if self._client and self._client.has_collection(temp_collection_name):
|
|
logger.info(
|
|
f"[{self.workspace}] Cleaning up failed migration temporary collection"
|
|
)
|
|
self._client.drop_collection(temp_collection_name)
|
|
except Exception as cleanup_error:
|
|
logger.warning(
|
|
f"[{self.workspace}] Failed to cleanup temporary collection: {cleanup_error}"
|
|
)
|
|
|
|
# Re-raise the original error
|
|
raise RuntimeError(
|
|
f"Iterator-based migration failed for collection {self.namespace}: {e}"
|
|
) from e
|
|
|
|
finally:
|
|
# Ensure iterator is properly closed
|
|
if iterator:
|
|
try:
|
|
iterator.close()
|
|
logger.debug(
|
|
f"[{self.workspace}] Query iterator closed successfully"
|
|
)
|
|
except Exception as close_error:
|
|
logger.warning(
|
|
f"[{self.workspace}] Failed to close query iterator: {close_error}"
|
|
)
|
|
|
|
def _validate_collection_compatibility(self):
|
|
"""Validate existing collection's dimension and schema compatibility"""
|
|
try:
|
|
collection_info = self._client.describe_collection(self.final_namespace)
|
|
|
|
# 1. Check vector dimension
|
|
self._check_vector_dimension(collection_info)
|
|
|
|
# 2. Check schema compatibility
|
|
self._check_schema_compatibility(collection_info)
|
|
|
|
logger.info(
|
|
f"[{self.workspace}] VectorDB Collection '{self.namespace}' compatibility validation passed"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Collection compatibility validation failed for {self.namespace}: {e}"
|
|
)
|
|
raise
|
|
|
|
@staticmethod
|
|
def _is_missing_vector_index_error(error: Exception) -> bool:
|
|
"""Return True when the error indicates the collection lacks a vector index."""
|
|
error_message = str(error).lower()
|
|
return (
|
|
"no vector index" in error_message
|
|
or "please create index firstly" in error_message
|
|
)
|
|
|
|
def _repair_missing_vector_index(self):
|
|
"""Create indexes for an existing collection that is missing its vector index."""
|
|
logger.warning(
|
|
f"[{self.workspace}] Collection '{self.namespace}' is missing a vector index, attempting repair"
|
|
)
|
|
self._create_indexes_after_collection()
|
|
|
|
def _ensure_collection_loaded(self):
|
|
"""Ensure the collection is loaded into memory for search operations"""
|
|
try:
|
|
# Check if collection exists first
|
|
if not self._client.has_collection(self.final_namespace):
|
|
logger.error(
|
|
f"[{self.workspace}] Collection {self.namespace} does not exist"
|
|
)
|
|
raise ValueError(f"Collection {self.final_namespace} does not exist")
|
|
|
|
# Load the collection if it's not already loaded
|
|
# In Milvus, collections need to be loaded before they can be searched
|
|
self._client.load_collection(self.final_namespace)
|
|
# logger.debug(f"[{self.workspace}] Collection {self.namespace} loaded successfully")
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Failed to load collection {self.namespace}: {e}"
|
|
)
|
|
raise
|
|
|
|
def _create_collection_if_not_exist(self):
|
|
"""Create collection if not exists and check existing collection compatibility"""
|
|
|
|
try:
|
|
# Check if our specific collection exists
|
|
collection_exists = self._client.has_collection(self.final_namespace)
|
|
logger.info(
|
|
f"[{self.workspace}] VectorDB collection '{self.namespace}' exists check: {collection_exists}"
|
|
)
|
|
|
|
if collection_exists:
|
|
# Double-check by trying to describe the collection
|
|
try:
|
|
self._client.describe_collection(self.final_namespace)
|
|
self._validate_collection_compatibility()
|
|
try:
|
|
# Ensure the collection is loaded after validation
|
|
self._ensure_collection_loaded()
|
|
return
|
|
except Exception as load_error:
|
|
if not self._is_missing_vector_index_error(load_error):
|
|
raise
|
|
|
|
try:
|
|
self._repair_missing_vector_index()
|
|
self._ensure_collection_loaded()
|
|
logger.info(
|
|
f"[{self.workspace}] Repaired missing vector index for existing collection '{self.namespace}'"
|
|
)
|
|
return
|
|
except Exception as repair_error:
|
|
raise RuntimeError(
|
|
f"Index repair failed for collection '{self.final_namespace}'. "
|
|
f"Original error: {repair_error}"
|
|
) from repair_error
|
|
except Exception as validation_error:
|
|
# CRITICAL: Collection exists but validation failed
|
|
# This indicates potential data migration failure or incompatible schema
|
|
# Stop execution to prevent data loss and require manual intervention
|
|
logger.error(
|
|
f"[{self.workspace}] CRITICAL ERROR: Collection '{self.namespace}' exists but validation failed!"
|
|
)
|
|
logger.error(
|
|
f"[{self.workspace}] This indicates potential data migration failure or schema incompatibility."
|
|
)
|
|
logger.error(
|
|
f"[{self.workspace}] Validation error: {validation_error}"
|
|
)
|
|
logger.error(f"[{self.workspace}] MANUAL INTERVENTION REQUIRED:")
|
|
logger.error(
|
|
f"[{self.workspace}] 1. Check the existing collection schema and data integrity"
|
|
)
|
|
logger.error(
|
|
f"[{self.workspace}] 2. Backup existing data if needed"
|
|
)
|
|
logger.error(
|
|
f"[{self.workspace}] 3. Manually resolve schema compatibility issues"
|
|
)
|
|
logger.error(
|
|
f"[{self.workspace}] 4. Consider dropping and recreating the collection if data is not critical"
|
|
)
|
|
logger.error(
|
|
f"[{self.workspace}] Program execution stopped to prevent potential data loss."
|
|
)
|
|
|
|
# Raise a specific exception to stop execution
|
|
raise RuntimeError(
|
|
f"Collection validation failed for '{self.final_namespace}'. "
|
|
f"Data migration failure detected. Manual intervention required to prevent data loss. "
|
|
f"Original error: {validation_error}"
|
|
)
|
|
|
|
# Collection doesn't exist, create new collection
|
|
logger.info(f"[{self.workspace}] Creating new collection: {self.namespace}")
|
|
schema = self._create_schema_for_namespace()
|
|
|
|
# Create collection with schema only first
|
|
self._client.create_collection(
|
|
collection_name=self.final_namespace, schema=schema
|
|
)
|
|
|
|
# Then create indexes
|
|
self._create_indexes_after_collection()
|
|
|
|
# Load the newly created collection
|
|
self._ensure_collection_loaded()
|
|
|
|
logger.info(
|
|
f"[{self.workspace}] Successfully created Milvus collection: {self.namespace}"
|
|
)
|
|
|
|
except RuntimeError:
|
|
# Re-raise RuntimeError (validation failures) without modification
|
|
# These are critical errors that should stop execution
|
|
raise
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error in _create_collection_if_not_exist for {self.namespace}: {e}"
|
|
)
|
|
|
|
# If there's any error (other than validation failure), try to force create the collection
|
|
logger.info(
|
|
f"[{self.workspace}] Attempting to force create collection {self.namespace}..."
|
|
)
|
|
try:
|
|
# Try to drop the collection first if it exists in a bad state
|
|
try:
|
|
if self._client.has_collection(self.final_namespace):
|
|
logger.info(
|
|
f"[{self.workspace}] Dropping potentially corrupted collection {self.namespace}"
|
|
)
|
|
self._client.drop_collection(self.final_namespace)
|
|
except Exception as drop_error:
|
|
logger.warning(
|
|
f"[{self.workspace}] Could not drop collection {self.namespace}: {drop_error}"
|
|
)
|
|
|
|
# Create fresh collection
|
|
schema = self._create_schema_for_namespace()
|
|
self._client.create_collection(
|
|
collection_name=self.final_namespace, schema=schema
|
|
)
|
|
self._create_indexes_after_collection()
|
|
|
|
# Load the newly created collection
|
|
self._ensure_collection_loaded()
|
|
|
|
logger.info(
|
|
f"[{self.workspace}] Successfully force-created collection {self.namespace}"
|
|
)
|
|
|
|
except Exception as create_error:
|
|
logger.error(
|
|
f"[{self.workspace}] Failed to force-create collection {self.namespace}: {create_error}"
|
|
)
|
|
raise
|
|
|
|
def __post_init__(self):
|
|
self._validate_embedding_func()
|
|
|
|
# Extract MilvusIndexConfig parameters from vector_db_storage_cls_kwargs
|
|
#
|
|
# IMPORTANT: This approach allows Milvus index configuration via vector_db_storage_cls_kwargs,
|
|
# which is the RECOMMENDED method for framework integration (e.g., RAGAnything).
|
|
#
|
|
# All 11 index configuration parameters can be passed through vector_db_storage_cls_kwargs:
|
|
# - index_type, metric_type
|
|
# - hnsw_m, hnsw_ef_construction, hnsw_ef
|
|
# - sq_type, sq_refine, sq_refine_type, sq_refine_k
|
|
# - ivf_nlist, ivf_nprobe
|
|
#
|
|
# Example:
|
|
# LightRAG(
|
|
# vector_storage="MilvusVectorDBStorage",
|
|
# vector_db_storage_cls_kwargs={
|
|
# "cosine_better_than_threshold": 0.2,
|
|
# "index_type": "HNSW",
|
|
# "metric_type": "COSINE",
|
|
# "hnsw_m": 32,
|
|
# "hnsw_ef_construction": 256,
|
|
# }
|
|
# )
|
|
#
|
|
# Use MilvusIndexConfig.get_config_field_names() to dynamically extract valid parameters.
|
|
# This ensures we always stay in sync with the MilvusIndexConfig dataclass definition.
|
|
kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
|
|
index_config_keys = MilvusIndexConfig.get_config_field_names()
|
|
index_config_params = {
|
|
k: v for k, v in kwargs.items() if k in index_config_keys
|
|
}
|
|
|
|
# Initialize index configuration (if not already set)
|
|
# Configuration priority: init params from kwargs > environment variables > defaults
|
|
if not hasattr(self, "index_config") or self.index_config is None:
|
|
self.index_config = MilvusIndexConfig(**index_config_params)
|
|
|
|
# Check for MILVUS_WORKSPACE environment variable first (higher priority)
|
|
# This allows administrators to force a specific workspace for all Milvus storage instances
|
|
milvus_workspace = os.environ.get("MILVUS_WORKSPACE")
|
|
if milvus_workspace and milvus_workspace.strip():
|
|
# Use environment variable value, overriding the passed workspace parameter
|
|
effective_workspace = milvus_workspace.strip()
|
|
logger.info(
|
|
f"Using MILVUS_WORKSPACE environment variable: '{effective_workspace}' (overriding '{self.workspace}/{self.namespace}')"
|
|
)
|
|
else:
|
|
# Use the workspace parameter passed during initialization
|
|
effective_workspace = self.workspace
|
|
if effective_workspace:
|
|
logger.debug(
|
|
f"Using passed workspace parameter: '{effective_workspace}'"
|
|
)
|
|
|
|
# Build final_namespace with workspace prefix for data isolation
|
|
# Keep original namespace unchanged for type detection logic
|
|
if effective_workspace:
|
|
self.final_namespace = f"{effective_workspace}_{self.namespace}"
|
|
logger.debug(
|
|
f"Final namespace with workspace prefix: '{self.final_namespace}'"
|
|
)
|
|
else:
|
|
# When workspace is empty, final_namespace equals original namespace
|
|
self.final_namespace = self.namespace
|
|
self.workspace = ""
|
|
logger.debug(f"Final namespace (no workspace): '{self.final_namespace}'")
|
|
cosine_threshold = kwargs.get("cosine_better_than_threshold")
|
|
if cosine_threshold is None:
|
|
raise ValueError(
|
|
"cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
|
|
)
|
|
self.cosine_better_than_threshold = cosine_threshold
|
|
|
|
# Ensure created_at is in meta_fields
|
|
if "created_at" not in self.meta_fields:
|
|
self.meta_fields.add("created_at")
|
|
|
|
# Initialize client as None - will be created in initialize() method
|
|
self._client = None
|
|
self._max_batch_size = self.global_config["embedding_batch_num"]
|
|
self._initialized = False
|
|
|
|
async def initialize(self):
|
|
"""Initialize Milvus collection"""
|
|
async with get_data_init_lock():
|
|
if self._initialized:
|
|
return
|
|
|
|
try:
|
|
# Create MilvusClient if not already created
|
|
if self._client is None:
|
|
self._client = self._create_milvus_client()
|
|
logger.debug(
|
|
f"[{self.workspace}] MilvusClient created successfully"
|
|
)
|
|
|
|
# Validate Milvus version compatibility with configured index
|
|
if self.index_config.index_type in INDEX_VERSION_REQUIREMENTS:
|
|
try:
|
|
server_version = self._client.get_server_version()
|
|
self.index_config.validate_milvus_version(server_version)
|
|
except Exception as version_error:
|
|
logger.error(
|
|
f"[{self.workspace}] Milvus version validation failed: {version_error}"
|
|
)
|
|
raise
|
|
|
|
# Create collection and check compatibility
|
|
self._create_collection_if_not_exist()
|
|
self._initialized = True
|
|
logger.info(
|
|
f"[{self.workspace}] Milvus collection '{self.namespace}' initialized successfully"
|
|
)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Failed to initialize Milvus collection '{self.namespace}': {e}"
|
|
)
|
|
raise
|
|
|
|
async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
|
|
# logger.debug(f"[{self.workspace}] Inserting {len(data)} to {self.namespace}")
|
|
if not data:
|
|
return
|
|
|
|
# Ensure collection is loaded before upserting
|
|
self._ensure_collection_loaded()
|
|
|
|
import time
|
|
|
|
current_time = int(time.time())
|
|
|
|
list_data: list[dict[str, Any]] = [
|
|
{
|
|
"id": k,
|
|
"created_at": current_time,
|
|
**{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
|
|
}
|
|
for k, v in data.items()
|
|
]
|
|
contents = [v["content"] for v in data.values()]
|
|
batches = [
|
|
contents[i : i + self._max_batch_size]
|
|
for i in range(0, len(contents), self._max_batch_size)
|
|
]
|
|
|
|
embedding_tasks = [
|
|
self.embedding_func(batch, context="document") for batch in batches
|
|
]
|
|
embeddings_list = await asyncio.gather(*embedding_tasks)
|
|
|
|
embeddings = np.concatenate(embeddings_list)
|
|
for i, d in enumerate(list_data):
|
|
d["vector"] = embeddings[i]
|
|
results = self._client.upsert(
|
|
collection_name=self.final_namespace, data=list_data
|
|
)
|
|
return results
|
|
|
|
async def query(
|
|
self, query: str, top_k: int, query_embedding: list[float] = None
|
|
) -> list[dict[str, Any]]:
|
|
# Ensure collection is loaded before querying
|
|
self._ensure_collection_loaded()
|
|
|
|
# Use provided embedding or compute it
|
|
if query_embedding is not None:
|
|
embedding = [query_embedding] # Milvus expects a list of embeddings
|
|
else:
|
|
embedding = await self.embedding_func(
|
|
[query], context="query", _priority=5
|
|
) # higher priority for query
|
|
|
|
# Include all meta_fields (created_at is now always included)
|
|
output_fields = list(self.meta_fields)
|
|
|
|
# Build search params from index config
|
|
search_params_base = self.index_config.build_search_params()
|
|
|
|
# Merge with metric type and radius threshold
|
|
search_params = {
|
|
"metric_type": self.index_config.metric_type,
|
|
"params": {
|
|
**search_params_base.get("params", {}),
|
|
"radius": self.cosine_better_than_threshold,
|
|
},
|
|
}
|
|
|
|
results = self._client.search(
|
|
collection_name=self.final_namespace,
|
|
data=embedding,
|
|
limit=top_k,
|
|
output_fields=output_fields,
|
|
search_params=search_params,
|
|
)
|
|
return [
|
|
{
|
|
**dp["entity"],
|
|
"id": dp["id"],
|
|
"distance": dp["distance"],
|
|
"created_at": dp.get("created_at"),
|
|
}
|
|
for dp in results[0]
|
|
]
|
|
|
|
async def index_done_callback(self) -> None:
|
|
# Milvus handles persistence automatically
|
|
pass
|
|
|
|
async def delete_entity(self, entity_name: str) -> None:
|
|
"""Delete an entity from the vector database
|
|
|
|
Args:
|
|
entity_name: The name of the entity to delete
|
|
"""
|
|
try:
|
|
# Compute entity ID from name
|
|
entity_id = compute_mdhash_id(entity_name, prefix="ent-")
|
|
logger.debug(
|
|
f"[{self.workspace}] Attempting to delete entity {entity_name} with ID {entity_id}"
|
|
)
|
|
|
|
# Delete the entity from Milvus collection
|
|
result = self._client.delete(
|
|
collection_name=self.final_namespace, pks=[entity_id]
|
|
)
|
|
|
|
if result and result.get("delete_count", 0) > 0:
|
|
logger.debug(
|
|
f"[{self.workspace}] Successfully deleted entity {entity_name}"
|
|
)
|
|
else:
|
|
logger.debug(
|
|
f"[{self.workspace}] Entity {entity_name} not found in storage"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"[{self.workspace}] Error deleting entity {entity_name}: {e}")
|
|
|
|
async def delete_entity_relation(self, entity_name: str) -> None:
|
|
"""Delete all relations associated with an entity
|
|
|
|
Args:
|
|
entity_name: The name of the entity whose relations should be deleted
|
|
"""
|
|
try:
|
|
# Ensure collection is loaded before querying
|
|
self._ensure_collection_loaded()
|
|
|
|
# Search for relations where entity is either source or target
|
|
expr = f'src_id == "{entity_name}" or tgt_id == "{entity_name}"'
|
|
|
|
# Find all relations involving this entity
|
|
results = self._client.query(
|
|
collection_name=self.final_namespace, filter=expr, output_fields=["id"]
|
|
)
|
|
|
|
if not results or len(results) == 0:
|
|
logger.debug(
|
|
f"[{self.workspace}] No relations found for entity {entity_name}"
|
|
)
|
|
return
|
|
|
|
# Extract IDs of relations to delete
|
|
relation_ids = [item["id"] for item in results]
|
|
logger.debug(
|
|
f"[{self.workspace}] Found {len(relation_ids)} relations for entity {entity_name}"
|
|
)
|
|
|
|
# Delete the relations
|
|
if relation_ids:
|
|
delete_result = self._client.delete(
|
|
collection_name=self.final_namespace, pks=relation_ids
|
|
)
|
|
|
|
logger.debug(
|
|
f"[{self.workspace}] Deleted {delete_result.get('delete_count', 0)} relations for {entity_name}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error deleting relations for {entity_name}: {e}"
|
|
)
|
|
|
|
async def delete(self, ids: list[str]) -> None:
|
|
"""Delete vectors with specified IDs
|
|
|
|
Args:
|
|
ids: List of vector IDs to be deleted
|
|
"""
|
|
try:
|
|
# Ensure collection is loaded before deleting
|
|
self._ensure_collection_loaded()
|
|
|
|
# Delete vectors by IDs
|
|
result = self._client.delete(collection_name=self.final_namespace, pks=ids)
|
|
|
|
if result and result.get("delete_count", 0) > 0:
|
|
logger.debug(
|
|
f"[{self.workspace}] Successfully deleted {result.get('delete_count', 0)} vectors from {self.namespace}"
|
|
)
|
|
else:
|
|
logger.debug(
|
|
f"[{self.workspace}] No vectors were deleted from {self.namespace}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error while deleting vectors from {self.namespace}: {e}"
|
|
)
|
|
|
|
async def get_by_id(self, id: str) -> dict[str, Any] | None:
|
|
"""Get vector data by its ID
|
|
|
|
Args:
|
|
id: The unique identifier of the vector
|
|
|
|
Returns:
|
|
The vector data if found, or None if not found
|
|
"""
|
|
try:
|
|
# Ensure collection is loaded before querying
|
|
self._ensure_collection_loaded()
|
|
|
|
# Include all meta_fields (created_at is now always included) plus id
|
|
output_fields = list(self.meta_fields) + ["id"]
|
|
|
|
# Query Milvus for a specific ID
|
|
result = self._client.query(
|
|
collection_name=self.final_namespace,
|
|
filter=f'id == "{id}"',
|
|
output_fields=output_fields,
|
|
)
|
|
|
|
if not result or len(result) == 0:
|
|
return None
|
|
|
|
return result[0]
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error retrieving vector data for ID {id}: {e}"
|
|
)
|
|
return None
|
|
|
|
async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
|
|
"""Get multiple vector data by their IDs
|
|
|
|
Args:
|
|
ids: List of unique identifiers
|
|
|
|
Returns:
|
|
List of vector data objects that were found
|
|
"""
|
|
if not ids:
|
|
return []
|
|
|
|
try:
|
|
# Ensure collection is loaded before querying
|
|
self._ensure_collection_loaded()
|
|
|
|
# Include all meta_fields (created_at is now always included) plus id
|
|
output_fields = list(self.meta_fields) + ["id"]
|
|
|
|
# Prepare the ID filter expression
|
|
id_list = '", "'.join(ids)
|
|
filter_expr = f'id in ["{id_list}"]'
|
|
|
|
# Query Milvus with the filter
|
|
result = self._client.query(
|
|
collection_name=self.final_namespace,
|
|
filter=filter_expr,
|
|
output_fields=output_fields,
|
|
)
|
|
|
|
if not result:
|
|
return []
|
|
|
|
result_map: dict[str, dict[str, Any]] = {}
|
|
for row in result:
|
|
if not row:
|
|
continue
|
|
row_id = row.get("id")
|
|
if row_id is not None:
|
|
result_map[str(row_id)] = row
|
|
|
|
ordered_results: list[dict[str, Any] | None] = []
|
|
for requested_id in ids:
|
|
ordered_results.append(result_map.get(str(requested_id)))
|
|
|
|
return ordered_results
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error retrieving vector data for IDs {ids}: {e}"
|
|
)
|
|
return []
|
|
|
|
async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
|
|
"""Get vectors by their IDs, returning only ID and vector data for efficiency
|
|
|
|
Args:
|
|
ids: List of unique identifiers
|
|
|
|
Returns:
|
|
Dictionary mapping IDs to their vector embeddings
|
|
Format: {id: [vector_values], ...}
|
|
"""
|
|
if not ids:
|
|
return {}
|
|
|
|
try:
|
|
# Ensure collection is loaded before querying
|
|
self._ensure_collection_loaded()
|
|
|
|
# Prepare the ID filter expression
|
|
id_list = '", "'.join(ids)
|
|
filter_expr = f'id in ["{id_list}"]'
|
|
|
|
# Query Milvus with the filter, requesting only vector field
|
|
result = self._client.query(
|
|
collection_name=self.final_namespace,
|
|
filter=filter_expr,
|
|
output_fields=["vector"],
|
|
)
|
|
|
|
vectors_dict = {}
|
|
for item in result:
|
|
if item and "vector" in item and "id" in item:
|
|
# Convert numpy array to list if needed
|
|
vector_data = item["vector"]
|
|
if isinstance(vector_data, np.ndarray):
|
|
vector_data = vector_data.tolist()
|
|
vectors_dict[item["id"]] = vector_data
|
|
|
|
return vectors_dict
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error retrieving vectors by IDs from {self.namespace}: {e}"
|
|
)
|
|
return {}
|
|
|
|
async def drop(self) -> dict[str, str]:
|
|
"""Drop all vector data from storage and clean up resources
|
|
|
|
This method will delete all data from the Milvus collection.
|
|
|
|
Returns:
|
|
dict[str, str]: Operation status and message
|
|
- On success: {"status": "success", "message": "data dropped"}
|
|
- On failure: {"status": "error", "message": "<error details>"}
|
|
"""
|
|
try:
|
|
# Drop the collection and recreate it
|
|
if self._client.has_collection(self.final_namespace):
|
|
self._client.drop_collection(self.final_namespace)
|
|
|
|
# Recreate the collection
|
|
self._create_collection_if_not_exist()
|
|
|
|
logger.info(
|
|
f"[{self.workspace}] Process {os.getpid()} drop Milvus collection {self.namespace}"
|
|
)
|
|
return {"status": "success", "message": "data dropped"}
|
|
except Exception as e:
|
|
logger.error(
|
|
f"[{self.workspace}] Error dropping Milvus collection {self.namespace}: {e}"
|
|
)
|
|
return {"status": "error", "message": str(e)}
|