isage-middleware 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

sage/middleware/components/neuromem/memory_collection/vdb_collection.py

@@ -0,0 +1,922 @@
+import os
+import json
+import yaml
+import shutil
+import inspect
+import numpy as np
+from typing import Optional, Dict, Any, List, Callable
+from sage.common.utils.logging.custom_logger import CustomLogger
+from sage.middleware.components.neuromem.memory_collection.base_collection import BaseMemoryCollection
+from sage.middleware.components.neuromem.search_engine.vdb_index import index_factory
+from sage.middleware.components.neuromem.utils.path_utils import get_default_data_dir
+from sage.middleware.utils.embedding.embedding_api import apply_embedding_model
+
+
+class VDBMemoryCollection(BaseMemoryCollection):
+    """
+    Memory collection with vector database support.
+    支持向量数据库功能的内存集合类。
+    
+    支持两种初始化方式：
+    1. 通过声明VDBMemoryCollection(config, corpus)创建
+    2. 通过VDBMemoryCollection.load(name, vdb_path)恢复式创建
+    """
+    def __init__(
+        self, 
+        config: Dict[str, Any]
+    ):
+        """
+        初始化VDBMemoryCollection
+        
+        Args:
+            config: 配置字典，必须包含name等参数
+        """
+        # 初始化CustomLogger
+        self.logger = CustomLogger()
+        
+        if "name" not in config:
+            self.logger.error("config中必须包含'name'字段")
+            raise ValueError("config中必须包含'name'字段")
+        
+        self.name = config["name"]
+        super().__init__(self.name)
+        
+        # 使用新的参数命名规范，删除向后兼容的代码
+        self.default_embedding_model_name = config.get("default_embedding_model", "default")
+        self.default_embedding_model = apply_embedding_model(self.default_embedding_model_name)
+        self.default_dim = config.get("default_dim", self.default_embedding_model.get_dim())
+        self.default_topk = config.get("default_topk", 5)
+        self.default_backend_type = config.get("default_vdb_backend", "FAISS")
+        
+        self.index_info = {}  # index_name -> dict: { embedding_model_name, dim, index, backend_type, topk, description, metadata_filter_func, metadata_conditions }
+        self.index_embedding_model = {} # index_name -> embedding_model 以供随时使用
+        
+        # 创建全局索引，优先使用config中的global_index配置，否则使用默认配置
+        global_index_config = config.get("global_index", {})
+        if "name" not in global_index_config:
+            global_index_config["name"] = "global_index"
+        if "embedding_model" not in global_index_config:
+            global_index_config["embedding_model"] = self.default_embedding_model_name
+        if "dim" not in global_index_config:
+            global_index_config["dim"] = self.default_dim
+        if "backend_type" not in global_index_config:
+            global_index_config["backend_type"] = self.default_backend_type
+        if "topk" not in global_index_config:
+            global_index_config["topk"] = self.default_topk
+        
+        self.create_index(config=global_index_config)
+    
+    def batch_insert_data(self, data: List[str], metadatas: Optional[List[Dict[str, Any]]] = None):
+        """
+        批量插入数据到collection中（仅存储，不创建索引）
+        
+        Args:
+            data: 数据列表（文本、图片等）
+            metadatas: 对应的元数据列表，可选
+        """
+        self.logger.info(f"Batch inserting {len(data)} data items to storage")
+        
+        if metadatas is not None and len(metadatas) != len(data):
+            raise ValueError("metadatas length must match data length")
+        
+        for i, item in enumerate(data):
+            stable_id = self._get_stable_id(item)
+            self.text_storage.store(stable_id, item)
+            
+            if metadatas and metadatas[i]:
+                metadata = metadatas[i]
+                # 自动注册所有未知的元数据字段
+                for field_name in metadata.keys():
+                    if not self.metadata_storage.has_field(field_name):
+                        self.metadata_storage.add_field(field_name)
+                self.metadata_storage.store(stable_id, metadata)
+    
+    def _serialize_func(self, func):
+        """
+        改善lambda序列化管理
+        """
+        if func is None:
+            return None
+        try:
+            return inspect.getsource(func).strip()
+        except Exception:
+            return str(func)
+    
+    def _deserialize_func(self, func_str):
+        """
+        反序列化函数字符串
+        """
+        if func_str is None or func_str == "None" or func_str == "":
+            return lambda m: True
+        
+        # 简单的lambda函数恢复，实际生产环境中需要更安全的方式
+        try:
+            # 这里只是一个简单的示例，实际应该使用更安全的方式
+            if func_str.startswith("lambda"):
+                return eval(func_str)
+            else:
+                return lambda m: True
+        except Exception:
+            return lambda m: True
+
+    def store(self, store_path: Optional[str] = None):
+        self.logger.debug(f"VDBMemoryCollection: store called")
+
+        if store_path is None:
+            # 使用默认数据目录
+            base_dir = get_default_data_dir()
+        else:
+            # 使用传入的数据目录（通常来自MemoryManager）
+            base_dir = store_path
+            
+        collection_dir = os.path.join(base_dir, "vdb_collection", self.name)
+        os.makedirs(collection_dir, exist_ok=True)
+
+        # 1. 存储text和metadata
+        self.text_storage.store_to_disk(os.path.join(collection_dir, "text_storage.json"))
+        self.metadata_storage.store_to_disk(os.path.join(collection_dir, "metadata_storage.json"))
+
+        # 2. 索引和index_info
+        indexes_dir = os.path.join(collection_dir, "indexes")
+        os.makedirs(indexes_dir, exist_ok=True)
+        saved_index_info = {}
+        for index_name, info in self.index_info.items():
+            idx = info["index"]
+            idx_path = os.path.join(indexes_dir, index_name)
+            os.makedirs(idx_path, exist_ok=True)
+            idx.store(idx_path)
+            saved_index_info[index_name] = {
+                "embedding_model_name": info.get("embedding_model_name", "default"),
+                "dim": info.get("dim", self.default_dim),
+                "backend_type": info.get("backend_type", "FAISS"),
+                "topk": info.get("topk", self.default_topk),
+                "description": info.get("description", ""),
+                "index_type": idx.__class__.__name__,
+                "metadata_filter_func": self._serialize_func(info.get("metadata_filter_func")),
+                "metadata_conditions": info.get("metadata_conditions", {}),
+            }
+
+        # 3. collection全局config
+        config = {
+            "name": self.name,
+            "default_embedding_model_name": self.default_embedding_model_name,
+            "default_dim": self.default_dim,
+            "default_topk": self.default_topk,
+            "default_backend_type": self.default_backend_type,
+            "indexes": saved_index_info,
+        }
+        with open(os.path.join(collection_dir, "config.json"), "w", encoding="utf-8") as f:
+            json.dump(config, f, ensure_ascii=False, indent=2)
+        
+        return {"collection_path": collection_dir}
+
+    @classmethod
+    def load(cls, name: str, vdb_path: Optional[str] = None):
+        """
+        从磁盘加载VDBMemoryCollection实例
+        
+        Args:
+            name: 集合名称
+            vdb_path: 加载路径，如果为None则使用默认路径
+        """
+        if vdb_path is None:
+            # 如果没有指定路径，使用默认路径结构
+            base_dir = get_default_data_dir()
+            load_path = os.path.join(base_dir, "vdb_collection", name)
+        else:
+            load_path = vdb_path
+        
+        # 此时 load_path 应该是指向具体collection的完整路径
+        config_path = os.path.join(load_path, "config.json")
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"No config found for collection at {config_path}")
+        
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+        
+        # 使用新的初始化方式创建实例
+        instance = cls(
+            config={
+                "name": name,
+                "default_embedding_model": config.get("default_embedding_model_name", config.get("embedding_model_name", "mockembedder")),
+                "default_dim": config.get("default_dim", config.get("dim", 128)),
+                "default_topk": config.get("default_topk", 5),
+                "default_vdb_backend": config.get("default_backend_type", config.get("backend_type", "FAISS"))
+            }
+        )
+        
+        # 加载storages (不再加载vector_storage)
+        instance.text_storage.load_from_disk(os.path.join(load_path, "text_storage.json"))
+        instance.metadata_storage.load_from_disk(os.path.join(load_path, "metadata_storage.json"))
+        
+        # 清空在初始化时创建的默认索引
+        instance.index_info.clear()
+        
+        # 加载索引和index_info
+        indexes_dir = os.path.join(load_path, "indexes")
+        for index_name, idx_info in config.get("indexes", {}).items():
+            idx_type = idx_info["index_type"]
+            idx_path = os.path.join(indexes_dir, index_name)
+            
+            try:
+                # 使用工厂类加载索引
+                if idx_type == "FaissIndex":
+                    idx = index_factory.load_index(index_name, "FAISS", idx_path)
+                else:
+                    # 尝试通过工厂类加载其他类型的索引
+                    backend_type = idx_type.replace("Index", "").upper()
+                    idx = index_factory.load_index(index_name, backend_type, idx_path)
+                    
+            except Exception as e:
+                raise NotImplementedError(f"Unknown index_type {idx_type}")
+            
+            # 恢复index_info
+            instance.index_info[index_name] = {
+                "embedding_model_name": idx_info.get("embedding_model_name", "default"),
+                "dim": idx_info.get("dim", config.get("default_dim", config.get("dim", 128))),
+                "index": idx,
+                "backend_type": idx_info.get("backend_type", "FAISS"),
+                "topk": idx_info.get("topk", 5),
+                "description": idx_info.get("description", ""),
+                "metadata_filter_func": instance._deserialize_func(idx_info.get("metadata_filter_func")),
+                "metadata_conditions": idx_info.get("metadata_conditions", {}),
+            }
+            
+            # 恢复embedding模型
+            embedding_model_name = idx_info.get("embedding_model_name", "default")
+            instance.index_embedding_model[index_name] = apply_embedding_model(embedding_model_name)
+
+        return instance
+
+    @staticmethod
+    def clear(name, clear_path=None):
+        if clear_path is None:
+            clear_path = get_default_data_dir()
+        collection_dir = os.path.join(clear_path, "vdb_collection", name)
+        try:
+            shutil.rmtree(collection_dir)
+            print(f"Cleared collection: {collection_dir}")
+        except FileNotFoundError:
+            print(f"Collection does not exist: {collection_dir}")
+        except Exception as e:
+            print(f"Failed to clear: {e}")
+    
+    def create_index(
+        self,
+        config: Optional[dict] = None,
+        metadata_filter_func: Optional[Callable[[Dict[str, Any]], bool]] = None,
+        **metadata_conditions
+    ):
+        """
+        使用元数据筛选条件创建新的向量索引。
+        """
+        # 检查1: config必须不为空且包含name字段
+        if config is None:
+            self.logger.warning("config不能为空，无法创建索引")
+            return 0
+            
+        if "name" not in config:
+            self.logger.warning("config中必须包含'name'字段，无法创建索引")
+            return 0
+            
+        index_name = config["name"]
+        
+        # 检查2: 如果索引已存在，不允许创建
+        if index_name in self.index_info:
+            self.logger.warning(f"索引 '{index_name}' 已存在，无法重复创建")
+            return 0
+        
+        # 从config中解包Collection级别的参数，如果没有则使用默认值
+        backend_type = config.get("backend_type", self.default_backend_type)
+        description = config.get("description", f"Index for {index_name}")
+        embedding_model_name = config.get("embedding_model", self.default_embedding_model_name)
+        dim = config.get("dim", self.default_dim)
+        topk = config.get("topk", self.default_topk)
+        
+        # 创建embedding模型
+        embedding_model = apply_embedding_model(embedding_model_name)
+        
+        # 准备传递给工厂的Index级别配置参数
+        # Collection负责参数检查和分离，工厂只负责创建索引
+        index_config = {
+            "name": index_name,
+            "dim": dim
+        }
+        # 将其他Index级别的配置参数传递给工厂
+        for key, value in config.items():
+            if key not in ["backend_type", "description", "embedding_model", "topk"]:
+                index_config[key] = value
+        
+        # 使用工厂类创建空索引，现在直接传递config
+        try:
+            # 使用新的基于config的创建方法，简化接口
+            index = index_factory.create_index_from_config(
+                config=index_config
+            )
+            
+            # 存储到index_info中
+            self.index_info[index_name] = {
+                "embedding_model_name": embedding_model_name,
+                "dim": dim,
+                "index": index,
+                "backend_type": backend_type,
+                "topk": topk,
+                "description": description,
+                "metadata_filter_func": metadata_filter_func,
+                "metadata_conditions": metadata_conditions,
+            }
+            
+            # 存储embedding模型到index_embedding_model中
+            self.index_embedding_model[index_name] = embedding_model
+            
+            self.logger.info(f"成功创建索引 '{index_name}'，后端类型: {backend_type}")
+            return 1  # 成功创建返回1
+            
+        except Exception as e:
+            self.logger.error(f"Failed to create index {index_name} with backend {backend_type}: {e}")
+            raise
+
+    # 直接删除某个索引
+    def delete_index(self, index_name: str):
+        """
+        删除指定名称的索引。
+        """
+        if index_name in self.index_info:
+            del self.index_info[index_name]
+        else:
+            raise ValueError(f"Index '{index_name}' does not exist.")
+            
+        # 同时删除embedding模型
+        if index_name in self.index_embedding_model:
+            del self.index_embedding_model[index_name]
+
+    # 列举索引信息
+    def list_index(self, *index_names) -> List[Dict[str, str]]:
+        """
+        列出指定的索引或所有索引及其描述信息。
+        
+        Args:
+            *index_names: 可选的索引名称，如果不指定则返回所有索引
+            
+        Returns:
+            List[Dict]: [{"name": ..., "description": ...}, ...]
+        """
+        if index_names:
+            # 如果指定了索引名称，只返回这些索引的信息
+            result = []
+            for name in index_names:
+                if name in self.index_info:
+                    result.append({"name": name, "description": self.index_info[name]["description"]})
+                else:
+                    self.logger.warning(f"索引 '{name}' 不存在")
+            return result
+        else:
+            # 如果没有指定，返回所有索引信息
+            return [
+                {"name": name, "description": info["description"]}
+                for name, info in self.index_info.items()
+            ]
+
+    # 单条文本插入（指定索引，否则全局）
+    def insert(
+        self,
+        raw_text: str,
+        metadata: Optional[Dict[str, Any]] = None,
+        *index_names
+    ) -> str:
+        self.logger.debug(f"VDBMemoryCollection: insert called")
+        
+        # 首先存储数据到storage
+        stable_id = self._get_stable_id(raw_text)
+        self.text_storage.store(stable_id, raw_text)
+
+        if metadata:
+            # 自动注册所有未知的元数据字段
+            for field_name in metadata.keys():
+                if not self.metadata_storage.has_field(field_name):
+                    self.metadata_storage.add_field(field_name)
+            self.metadata_storage.store(stable_id, metadata)
+
+        # 如果没有指定索引，插入到全局索引
+        if not index_names:
+            if "global_index" not in self.index_info:
+                # 如果全局索引不存在，创建它
+                self.logger.info(f"创建全局索引: global_index")
+                self.create_index(config={"name": "global_index"})
+            index_names = ("global_index",)
+        
+        # 检查指定的索引是否存在
+        index_names_set = set(index_names)
+        for index_name in index_names_set:
+            if index_name not in self.index_info:
+                self.logger.warning(f"指定的索引 '{index_name}' 不存在，插入操作失败")
+                return "0"  # 返回0表示失败
+
+        # 修正插入逻辑：为每个索引使用对应的embedding模型
+        for index_name in index_names_set:
+            # 获取该索引对应的embedding模型
+            if index_name in self.index_embedding_model:
+                embedding_model = self.index_embedding_model[index_name]
+            else:
+                # 如果索引没有对应的embedding模型，使用默认模型
+                self.logger.warning(f"索引 '{index_name}' 没有对应的embedding模型，使用默认模型")
+                embedding_model = self.default_embedding_model
+            
+            # 使用对应的embedding模型编码文本
+            embedding = embedding_model.encode(raw_text)
+            
+            # 统一处理不同格式的embedding结果
+            if hasattr(embedding, "detach") and hasattr(embedding, "cpu"):
+                # PyTorch tensor
+                embedding = embedding.detach().cpu().numpy()
+            if isinstance(embedding, list):
+                # Python list
+                embedding = np.array(embedding)
+            if not isinstance(embedding, np.ndarray):
+                # 其他类型，尝试转换为numpy数组
+                embedding = np.array(embedding)
+                
+            # 确保数据类型是float32
+            embedding = embedding.astype(np.float32)
+            
+            # 检查embedding维度是否与索引要求一致
+            expected_dim = self.index_info[index_name]["dim"]
+            if embedding.shape[-1] != expected_dim:
+                self.logger.warning(f"索引 '{index_name}' 要求维度 {expected_dim}，但embedding维度为 {embedding.shape[-1]}，跳过插入")
+                continue
+            
+            # 插入到对应的索引中
+            index = self.index_info[index_name]["index"]
+            index.insert(embedding, stable_id)
+
+        return stable_id
+
+    # 单条文本更新（指定索引更新）
+    def update(
+        self,
+        former_data: str,
+        new_data: str,
+        new_metadata: Optional[Dict[str, Any]] = None,
+        *index_names: str
+    ) -> str:
+        old_id = self._get_stable_id(former_data)
+        if not self.text_storage.has(old_id):
+            raise ValueError("Original data not found.")
+
+        self.text_storage.delete(old_id)
+        self.metadata_storage.delete(old_id)
+
+        for index_info in self.index_info.values():
+            index_info["index"].delete(old_id)
+
+        return self.insert(new_data, new_metadata, *index_names)
+
+    # 单条文本删除（全索引删除）
+    def delete(self, raw_text: str):
+        stable_id = self._get_stable_id(raw_text)
+        self.text_storage.delete(stable_id)
+        self.metadata_storage.delete(stable_id)
+
+        for index_info in self.index_info.values():
+            index_info["index"].delete(stable_id)
+
+    def retrieve(
+        self,
+        raw_data: str,
+        topk: Optional[int] = None,
+        index_name: Optional[str] = None,
+        threshold: Optional[float] = None,
+        with_metadata: bool = False,
+        metadata_filter_func: Optional[Callable[[Dict[str, Any]], bool]] = None,
+        **metadata_conditions
+    ):
+        self.logger.debug(f"VDBMemoryCollection: retrieve called")   
+        
+        # 如果没有指定索引，使用或创建全局索引
+        if index_name is None:
+            index_name = "global_index"
+            if index_name not in self.index_info:
+                self.logger.info(f"Creating global index: {index_name}")
+                # 创建全局索引时，需要考虑元数据过滤条件
+                config = {"name": index_name}
+                self.create_index(config=config, metadata_filter_func=metadata_filter_func, **metadata_conditions)
+
+        if index_name not in self.index_info:
+            raise ValueError(f"Index '{index_name}' does not exist.")
+
+        if topk is None:
+            topk = int(self.default_topk)
+
+        # 使用对应索引的embedding模型
+        if index_name in self.index_embedding_model:
+            embedding_model = self.index_embedding_model[index_name]
+        else:
+            self.logger.warning(f"索引 '{index_name}' 没有对应的embedding模型，使用默认模型")
+            embedding_model = self.default_embedding_model
+
+        query_embedding = embedding_model.encode(raw_data)
+
+        # 统一处理不同格式的embedding结果
+        if hasattr(query_embedding, "detach") and hasattr(query_embedding, "cpu"):
+            # PyTorch tensor
+            query_embedding = query_embedding.detach().cpu().numpy()
+        if isinstance(query_embedding, list):
+            # Python list
+            query_embedding = np.array(query_embedding)
+        if not isinstance(query_embedding, np.ndarray):
+            # 其他类型，尝试转换为numpy数组
+            query_embedding = np.array(query_embedding)
+            
+        # 确保数据类型是float32
+        query_embedding = query_embedding.astype(np.float32)
+            
+        sub_index = self.index_info[index_name]["index"]
+        # 增加检索数量以补偿过滤后可能的损失
+        search_topk = topk * 2  # 检索更多结果以确保过滤后有足够的结果
+        top_k_ids, distances = sub_index.search(query_embedding, topk=search_topk, threshold=threshold)
+
+        if top_k_ids and isinstance(top_k_ids[0], (list, np.ndarray)):
+            top_k_ids = top_k_ids[0]
+        if distances and isinstance(distances[0], (list, np.ndarray)):
+            distances = distances[0]
+        top_k_ids = [str(i) for i in top_k_ids]
+
+        # 应用元数据过滤
+        if metadata_filter_func or metadata_conditions:
+            filtered_ids = self.filter_ids(top_k_ids, metadata_filter_func, **metadata_conditions)
+        else:
+            filtered_ids = top_k_ids
+
+        # 截取需要的数量，检索到几个就返回几个
+        final_ids = filtered_ids[:topk]
+
+        # 如果检索结果少于请求数量，记录信息但不警告
+        if len(final_ids) < topk:
+            self.logger.info(f"Retrieved {len(final_ids)} results (requested {topk})")
+
+        if with_metadata:
+            return [{"text": self.text_storage.get(i), "metadata": self.metadata_storage.get(i)} for i in final_ids]
+        else:
+            return [self.text_storage.get(i) for i in final_ids]
+
+    def update_index(self, index_name: str):
+        """
+        更新指定索引：删除当前索引，保留config，重新创建索引并批量插入数据
+        
+        Args:
+            index_name: 要更新的索引名称
+        """
+        if index_name not in self.index_info:
+            raise ValueError(f"Index '{index_name}' does not exist.")
+        
+        # 保存原始配置信息
+        info = self.index_info[index_name]
+        original_config = {
+            "name": index_name,
+            "backend_type": info["backend_type"],
+            "description": info["description"],
+            "embedding_model": info["embedding_model_name"],
+            "dim": info["dim"],
+            "topk": info["topk"]
+        }
+        original_metadata_filter_func = info.get("metadata_filter_func")
+        original_metadata_conditions = info.get("metadata_conditions", {})
+        
+        self.logger.info(f"开始更新索引: {index_name}")
+        
+        # 删除当前索引
+        self.delete_index(index_name)
+        
+        # 重新创建索引
+        self.create_index(
+            config=original_config,
+            metadata_filter_func=original_metadata_filter_func,
+            **original_metadata_conditions
+        )
+        
+        # 从storage中获取所有数据
+        all_ids = self.get_all_ids()
+        
+        if not all_ids:
+            self.logger.warning("storage中没有数据，索引更新完成但为空")
+            return
+        
+        # 应用元数据过滤（如果index_info中有规则）
+        filtered_ids = []
+        if original_metadata_filter_func or original_metadata_conditions:
+            for item_id in all_ids:
+                metadata = self.metadata_storage.get(item_id)
+                # 应用过滤函数
+                if original_metadata_filter_func and not original_metadata_filter_func(metadata or {}):
+                    continue
+                # 应用条件过滤
+                if original_metadata_conditions:
+                    match = True
+                    for key, value in original_metadata_conditions.items():
+                        if not metadata or metadata.get(key) != value:
+                            match = False
+                            break
+                    if not match:
+                        continue
+                filtered_ids.append(item_id)
+        else:
+            filtered_ids = all_ids
+        
+        # 获取对应的embedding模型
+        embedding_model = self.index_embedding_model[index_name]
+        
+        # 批量编码和插入
+        vectors = []
+        valid_ids = []
+        
+        for item_id in filtered_ids:
+            text = self.text_storage.get(item_id)
+            if text:
+                embedding = embedding_model.encode(text)
+                
+                # 统一处理不同格式的embedding结果
+                if hasattr(embedding, "detach") and hasattr(embedding, "cpu"):
+                    embedding = embedding.detach().cpu().numpy()
+                if isinstance(embedding, list):
+                    embedding = np.array(embedding)
+                if not isinstance(embedding, np.ndarray):
+                    embedding = np.array(embedding)
+                embedding = embedding.astype(np.float32)
+                
+                vectors.append(embedding)
+                valid_ids.append(item_id)
+        
+        # 使用底层的batch_insert方法批量插入
+        if vectors and valid_ids:
+            index = self.index_info[index_name]["index"]
+            result = index.batch_insert(vectors, valid_ids)
+            self.logger.info(f"索引 '{index_name}' 更新完成，插入了 {result} 条数据")
+        else:
+            self.logger.warning(f"索引 '{index_name}' 更新完成，但没有找到符合条件的数据")
+
+if __name__ == "__main__":
+    import torch
+    from transformers import AutoTokenizer, AutoModel
+    import shutil
+    import tempfile
+
+    def colored(text, color):
+        colors = {"green": "\033[92m", "red": "\033[91m", "yellow": "\033[93m", "reset": "\033[0m"}
+        return colors.get(color, "") + str(text) + colors["reset"]
+
+    class MockEmbeddingModel:
+        def encode(self, text):
+            # 模拟embedding，将文本长度作为特征
+            return torch.tensor([float(len(text))] * 4)  # 4维向量
+
+    def run_test():
+        print(colored("\n=== 开始VDBMemoryCollection测试 ===", "yellow"))
+        
+        # 准备测试环境
+        test_name = "test_collection"
+        test_dir = tempfile.mkdtemp()
+
+        try:
+            # 1. 测试新的初始化方式
+            print(colored("\n1. 测试新的初始化方式", "yellow"))
+            
+            # 方式1：通过config创建，使用mockembedder确保维度一致
+            config = {
+                "name": test_name,
+                "default_embedding_model": "mockembedder",
+                "default_dim": 128,  # 与mockembedder的固定维度一致
+                "default_topk": 5,
+                "default_vdb_backend": "FAISS"
+            }
+            collection = VDBMemoryCollection(config=config)
+            print(colored("✓ 通过config初始化成功", "green"))
+            
+            # 方式2：测试batch_insert_data
+            corpus = ["第一条文本", "第二条文本", "第三条文本"]
+            config_with_corpus = {
+                "name": f"{test_name}_corpus",
+                "default_embedding_model": "mockembedder",
+                "default_dim": 128,
+                "default_topk": 5,
+                "default_vdb_backend": "FAISS"
+            }
+            collection_with_corpus = VDBMemoryCollection(config=config_with_corpus)
+            collection_with_corpus.batch_insert_data(corpus)
+            # batch_insert_data 只存储数据，需要手动插入到索引
+            for text in corpus:
+                collection_with_corpus.insert(text)
+            print(colored("✓ 通过batch_insert_data成功", "green"))
+
+            # 2. 测试插入
+            print(colored("\n2. 测试数据插入", "yellow"))
+            texts = [
+                "这是第一条测试文本",
+                "这是第二条测试文本，带有metadata",
+                "这是第三条测试文本"
+            ]
+            metadata = {"type": "test", "priority": "high"}
+            
+            # 插入文本，不指定索引（应该会使用global_index）
+            id1 = collection.insert(texts[0])
+            # 插入文本，带metadata
+            id2 = collection.insert(texts[1], metadata=metadata)
+            # 插入文本到指定索引
+            collection.create_index(config={"name": "custom_index", "description": "自定义测试索引"})
+            id3 = collection.insert(texts[2], None, "custom_index")
+            
+            print(colored("✓ 数据插入成功", "green"))
+
+            # 3. 测试检索
+            print(colored("\n3. 测试检索功能", "yellow"))
+            
+            # 测试全局索引检索
+            results = collection.retrieve("测试文本", topk=2)
+            print(f"全局索引检索结果数量: {len(results)}")
+            assert len(results) > 0, "全局索引检索失败"
+            
+            # 测试指定索引检索
+            results = collection.retrieve("测试文本", topk=2, index_name="custom_index")
+            print(f"自定义索引检索结果数量: {len(results)}")
+            assert len(results) > 0, "自定义索引检索失败"
+            
+            # 测试带metadata的检索
+            results = collection.retrieve(
+                "测试文本",
+                topk=2,
+                with_metadata=True,
+                metadata_filter_func=lambda m: m and m.get("priority") == "high"
+            )
+            assert any(r.get("metadata", {}).get("priority") == "high" for r in results if isinstance(r, dict)), "metadata过滤失败"
+            
+            print(colored("✓ 检索功能测试通过", "green"))
+
+            # 4. 测试更新和删除
+            print(colored("\n4. 测试更新和删除", "yellow"))
+            
+            # 测试更新
+            new_text = "这是更新后的文本"
+            collection.update(texts[0], new_text)
+            results = collection.retrieve(new_text, topk=1)
+            assert results[0] == new_text, "更新操作失败"
+            
+            # 测试删除
+            collection.delete(texts[1])
+            
+            print(colored("✓ 更新和删除功能测试通过", "green"))
+
+            # 5. 测试持久化
+            print(colored("\n5. 测试持久化", "yellow"))
+            
+            # 保存
+            save_path = os.path.join(test_dir, "save_test")
+            collection.store(save_path)
+            
+            # 测试新的load方式
+            collection_dir = os.path.join(save_path, "vdb_collection", test_name)
+            loaded_collection = VDBMemoryCollection.load(test_name, collection_dir)
+            results = loaded_collection.retrieve("测试文本", topk=1)
+            assert len(results) > 0, "持久化后检索失败"
+            
+            print(colored("✓ 持久化功能测试通过", "green"))
+
+            # 6. 测试batch_insert_data功能
+            print(colored("\n6. 测试batch_insert_data功能", "yellow"))
+            corpus_results = collection_with_corpus.retrieve("文本", topk=3)
+            print(f"从batch_insert_data的集合检索结果数量: {len(corpus_results)}")
+            assert len(corpus_results) > 0, "batch_insert_data集合检索失败"
+            print(colored("✓ batch_insert_data功能测试通过", "green"))
+
+            # 7. 测试update_index功能
+            print(colored("\n7. 测试update_index功能", "yellow"))
+            # 向collection中添加更多数据
+            collection.insert("测试update_index的文本1")
+            collection.insert("测试update_index的文本2")
+            # 更新global_index
+            collection.update_index("global_index")
+            results = collection.retrieve("update_index", topk=5)
+            print(f"update_index后检索结果数量: {len(results)}")
+            print(colored("✓ update_index功能测试通过", "green"))
+
+            # 8. 测试lambda函数在持久化后的工作情况
+            print(colored("\n8. 测试lambda函数持久化", "yellow"))
+            
+            # 创建一个新的collection专门测试lambda
+            lambda_test_name = "lambda_test_collection"
+            lambda_config = {
+                "name": lambda_test_name,
+                "default_embedding_model": "mockembedder",
+                "default_dim": 128,
+                "default_topk": 5,
+                "default_vdb_backend": "FAISS"
+            }
+            lambda_collection = VDBMemoryCollection(config=lambda_config)
+            
+            # 添加测试数据
+            test_data = [
+                ("重要文档1", {"priority": "high", "category": "important"}),
+                ("普通文档2", {"priority": "low", "category": "normal"}),
+                ("重要文档3", {"priority": "high", "category": "important"}),
+                ("普通文档4", {"priority": "medium", "category": "normal"}),
+            ]
+            
+            for text, metadata in test_data:
+                lambda_collection.insert(text, metadata)
+            
+            # 创建一个带有lambda函数的索引
+            high_priority_filter = lambda m: m and m.get("priority") == "high"
+            lambda_collection.create_index(
+                config={
+                    "name": "high_priority_index",
+                    "description": "只包含高优先级文档的索引"
+                },
+                metadata_filter_func=high_priority_filter
+            )
+            
+            # 向带lambda的索引插入数据
+            for text, metadata in test_data:
+                if high_priority_filter(metadata):
+                    lambda_collection.insert(text, metadata, "high_priority_index")
+            
+            # 测试保存前的lambda过滤效果
+            print("保存前测试lambda过滤...")
+            before_save_results = lambda_collection.retrieve(
+                "文档", 
+                topk=10, 
+                index_name="high_priority_index"
+            )
+            print(f"保存前高优先级索引检索到 {len(before_save_results)} 条结果")
+            assert len(before_save_results) == 2, f"预期2条高优先级结果，实际得到{len(before_save_results)}条"
+            
+            # 保存collection
+            lambda_save_path = os.path.join(test_dir, "lambda_save_test")
+            lambda_collection.store(lambda_save_path)
+            print("lambda collection保存完成")
+            
+            # 加载collection
+            lambda_collection_dir = os.path.join(lambda_save_path, "vdb_collection", lambda_test_name)
+            loaded_lambda_collection = VDBMemoryCollection.load(lambda_test_name, lambda_collection_dir)
+            print("lambda collection加载完成")
+            
+            # 测试加载后的lambda过滤效果
+            print("加载后测试lambda过滤...")
+            after_load_results = loaded_lambda_collection.retrieve(
+                "文档", 
+                topk=10, 
+                index_name="high_priority_index"
+            )
+            print(f"加载后高优先级索引检索到 {len(after_load_results)} 条结果")
+            
+            # 验证结果是否一致
+            assert len(after_load_results) == len(before_save_results), \
+                f"加载后结果数量不一致：保存前{len(before_save_results)}条，加载后{len(after_load_results)}条"
+            
+            # 测试lambda函数在检索时的工作情况
+            print("测试加载后lambda函数在检索时的过滤效果...")
+            filtered_results = loaded_lambda_collection.retrieve(
+                "文档",
+                topk=10,
+                with_metadata=True,
+                metadata_filter_func=lambda m: m and m.get("priority") == "high"
+            )
+            
+            high_priority_count = sum(1 for r in filtered_results 
+                                    if isinstance(r, dict) and 
+                                       r.get("metadata", {}).get("priority") == "high")
+            print(f"通过lambda过滤检索到 {high_priority_count} 条高优先级结果")
+            assert high_priority_count > 0, "lambda函数在加载后没有正确工作"
+            
+            # 测试update_index是否能正确处理lambda函数
+            print("测试update_index对lambda函数的处理...")
+            loaded_lambda_collection.insert("新增重要文档", {"priority": "high", "category": "new"})
+            loaded_lambda_collection.update_index("high_priority_index")
+            
+            updated_results = loaded_lambda_collection.retrieve(
+                "文档", 
+                topk=10, 
+                index_name="high_priority_index"
+            )
+            print(f"update_index后高优先级索引检索到 {len(updated_results)} 条结果")
+            # 应该包含新增的重要文档，所以结果应该增加
+            assert len(updated_results) >= len(after_load_results), \
+                "update_index后lambda函数没有正确处理新数据"
+            
+            print(colored("✓ lambda函数持久化测试通过", "green"))
+            print("  - lambda函数序列化/反序列化正常")
+            print("  - 持久化后检索过滤功能正常")
+            print("  - update_index正确处理lambda函数")
+
+            print(colored("\n=== 所有测试通过！===", "green"))
+
+        except Exception as e:
+            print(colored(f"\n测试失败: {str(e)}", "red"))
+            import traceback
+            traceback.print_exc()
+            raise
+        finally:
+            # 清理测试数据
+            try:
+                shutil.rmtree(test_dir)
+            except:
+                pass
+
+    run_test()