isage-middleware 0.1.3.1__py3-none-any.whl → 0.2.0.0__py3-none-any.whl

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

@@ -1,709 +0,0 @@
-import os
-import json
-import yaml
-import shutil
-import inspect
-import warnings
-
-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.kv_index import KVIndexFactory
-from sage.middleware.components.neuromem.utils.path_utils import get_default_data_dir
-
-# 通过config文件指定默认索引，neuromem默认索引，用户指定索引
-
-def load_config(path: str) -> dict:
-    """加载YAML配置文件"""
-    with open(path, 'r') as f:
-        return yaml.safe_load(f)
-
-class KVMemoryCollection(BaseMemoryCollection):
-    """
-    基于键值对的内存集合，继承自 BaseMemoryCollection
-    提供基本的键值存储和检索功能
-    
-    支持两种初始化方式：
-    1. 通过config字典创建：KVMemoryCollection(config)
-    2. 通过load方法恢复：KVMemoryCollection.load(name, load_path)
-    """
-    def __init__(
-        self,
-        config: Dict[str, Any]
-    ):
-        """
-        初始化KVMemoryCollection
-        
-        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)
-        
-        # 从config中获取配置参数，提供默认值
-        self.default_topk = config.get("default_topk", 5)
-        self.default_index_type = config.get("default_index_type", "bm25s")
-        
-        self.indexes = {} # index_name -> {index_type, description, metadata_filter_func, metadata_conditions}
-        
-        # 如果config中指定了config_path，加载外部配置
-        config_path = config.get("config_path")
-        if config_path is not None:
-            external_config = load_config(config_path)
-            self.default_topk = external_config.get("kv_default_topk", self.default_topk)
-            self.default_index_type = external_config.get("kv_default_index_type", self.default_index_type)
-        
-        self.logger.info(f"KVMemoryCollection '{self.name}' 初始化成功")
-
-    def _serialize_func(self, func):
-        """
-        序列化函数以便持久化存储
-        """
-        if func is None:
-            return None
-        try:
-            return inspect.getsource(func).strip()
-        except Exception:
-            return str(func)  # 对于lambda等可能只能保存 repr
-
-    def _deserialize_func(self, func_str):
-        """
-        反序列化函数字符串
-        """
-        if func_str is None or func_str == "None" or func_str == "":
-            return None
-        
-        # 简单的lambda函数恢复，实际生产环境中需要更安全的方式
-        try:
-            # 这里只是一个简单的示例，实际应该使用更安全的方式
-            if func_str.startswith("lambda"):
-                return eval(func_str)
-            else:
-                # 对于其他函数类型，返回None，让调用者处理
-                return None
-        except Exception:
-            self.logger.warning(f"无法反序列化函数: {func_str}")
-            return None
-
-    @classmethod
-    def load(cls, name: str, load_path: Optional[str] = None) -> "KVMemoryCollection":
-        """
-        从磁盘加载KVMemoryCollection实例
-        
-        Args:
-            name: 集合名称
-            load_path: 加载路径，如果为None则使用默认路径
-        """
-        if load_path is None:
-            load_path = os.path.join(get_default_data_dir(), "kv_collection", name)
-        else:
-            # 当传入了具体路径时，直接使用该路径，不再添加额外层级
-            load_path = load_path
-        
-        # 创建实例时使用新的config方式
-        config = {"name": name}
-        instance = cls(config)
-        
-        # 加载数据
-        instance._load(load_path)
-        
-        return instance
-
-
-    def store(self, store_path: Optional[str] = None) -> Dict[str, Any]:
-        """
-        将集合保存到磁盘
-        
-        Args:
-            store_path: 保存路径，如果为None则使用默认路径
-        """
-        self.logger.debug(f"KVMemoryCollection: store called")
-        
-        if store_path is None:
-            store_path = get_default_data_dir()
-        # 加上kv_collection
-        collection_dir = os.path.join(store_path, "kv_collection", self.name)
-        os.makedirs(collection_dir, exist_ok=True)
-
-        # 存储 text 和 metadata
-        text_path = os.path.join(collection_dir, "text_storage.json")
-        metadata_path = os.path.join(collection_dir, "metadata_storage.json")
-        self.text_storage.store_to_disk(text_path)
-        self.metadata_storage.store_to_disk(metadata_path)
-
-        # 存储每个 index
-        index_info = {}
-        for index_name, info in self.indexes.items():
-            idx_type = info["index_type"]
-            idx = info["index"]
-            idx_type_dir = os.path.join(collection_dir, idx_type)
-            idx_path = os.path.join(idx_type_dir, index_name)
-            os.makedirs(idx_path, exist_ok=True)
-            idx.store(idx_path)
-            index_info[index_name] = {
-                "index_type": idx_type,
-                "description": info.get("description", ""),
-                "metadata_filter_func": self._serialize_func(info.get("metadata_filter_func")),
-                "metadata_conditions": info.get("metadata_conditions", {}),
-            }
-
-        config = {
-            "name": self.name,
-            "default_topk": self.default_topk,
-            "default_index_type": self.default_index_type,
-            "indexes": index_info,
-        }
-        config_path = os.path.join(collection_dir, "config.json")
-        with open(config_path, "w", encoding="utf-8") as f:
-            json.dump(config, f, ensure_ascii=False, indent=2)
-        
-        self.logger.info(f"集合 '{self.name}' 保存成功，路径: {collection_dir}")
-        return {"collection_path": collection_dir}
-
-    def _load(self, load_path: str):
-        """从指定路径加载集合数据"""
-        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)
-        self.default_topk = config.get("default_topk", 5)
-        self.default_index_type = config.get("default_index_type", "bm25s")
-
-        # 恢复 text 和 metadata
-        text_path = os.path.join(load_path, "text_storage.json")
-        metadata_path = os.path.join(load_path, "metadata_storage.json")
-        self.text_storage.load_from_disk(text_path)
-        self.metadata_storage.load_from_disk(metadata_path)
-
-        # 加载各 index
-        for index_name, idx_info in config.get("indexes", {}).items():
-            idx_type = idx_info["index_type"]
-            idx_path = os.path.join(load_path, idx_type, index_name)
-            
-            try:
-                idx = KVIndexFactory.load_index(idx_type, index_name, idx_path)
-            except ValueError as e:
-                raise NotImplementedError(f"Index type {idx_type} not supported: {e}")
-                
-            self.indexes[index_name] = {
-                "index": idx,
-                "index_type": idx_type,
-                "description": idx_info.get("description", ""),
-                "metadata_filter_func": self._deserialize_func(idx_info.get("metadata_filter_func")),
-                "metadata_conditions": idx_info.get("metadata_conditions", {}),
-            }
-        
-        self.logger.info(f"成功加载集合 '{self.name}'，包含 {len(self.indexes)} 个索引")
-
-    @staticmethod
-    def clear(name: str, clear_path: Optional[str] = None) -> None:
-        """
-        清理指定的集合
-        
-        Args:
-            name: 集合名称
-            clear_path: 清理路径，如果为None则使用默认路径
-        """
-        logger = CustomLogger()
-        
-        if clear_path is None:
-            clear_path = get_default_data_dir()
-        collection_dir = os.path.join(clear_path, "kv_collection", name)
-        try:
-            shutil.rmtree(collection_dir)
-            logger.info(f"成功清理集合: {collection_dir}")
-        except FileNotFoundError:
-            logger.warning(f"集合不存在: {collection_dir}")
-        except Exception as e:
-            logger.error(f"清理失败: {e}")
-            raise
-
-    def insert(
-        self,
-        raw_text: str, 
-        metadata: Optional[Dict[str, Any]] = None,
-        *index_names: str
-    ):
-        """
-        插入文本数据到指定索引
-        
-        Args:
-            raw_text: 原始文本
-            metadata: 元数据字典
-            *index_names: 要插入的索引名称列表
-        """
-        self.logger.debug(f"KVMemoryCollection: insert called")
-        
-        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)
-                
-        for index_name in index_names:
-            if index_name not in self.indexes:
-                self.logger.warning(f"Index '{index_name}' does not exist.")
-                continue
-            index = self.indexes[index_name]["index"]
-            index.insert(raw_text, stable_id)
-        
-        self.logger.debug(f"成功插入文本到 {len(index_names)} 个索引中")        
-        return stable_id
-    
-    def delete(self, raw_text: str):
-        """
-        删除指定文本及其关联数据
-        
-        Args:
-            raw_text: 要删除的原始文本
-        """
-        self.logger.debug(f"KVMemoryCollection: delete called")
-        
-        stable_id = self._get_stable_id(raw_text)
-        self.text_storage.delete(stable_id)
-        self.metadata_storage.delete(stable_id)
-
-        for index in self.indexes.values():
-            index["index"].delete(stable_id)
-        
-        self.logger.debug(f"成功删除文本，ID: {stable_id[:8]}...")
-    
-    def update( 
-        self,
-        former_text: str,
-        new_text: str,
-        new_metadata: Optional[Dict[str, Any]] = None,
-        *index_names: str
-    ) -> str:
-        """
-        更新文本数据
-        
-        Args:
-            former_text: 原始文本
-            new_text: 新文本
-            new_metadata: 新的元数据
-            *index_names: 要更新的索引名称列表
-        """
-        self.logger.debug(f"KVMemoryCollection: update called")
-        
-        old_id = self._get_stable_id(former_text)
-        if not self.text_storage.has(old_id):
-            raise ValueError("Original text not found.")
-
-        self.text_storage.delete(old_id)
-        self.metadata_storage.delete(old_id)
-
-        for index in self.indexes.values():
-            index["index"].delete(old_id)
-            
-        new_id = self.insert(new_text, new_metadata, *index_names)
-        self.logger.debug(f"成功更新文本，旧ID: {old_id[:8]}..., 新ID: {new_id[:8]}...")
-        return new_id
-    
-    def retrieve(
-        self,
-        raw_text: str,
-        topk: Optional[int] = None,
-        with_metadata: Optional[bool] = False,
-        index_name: Optional[str] = None,
-        metadata_filter_func: Optional[Callable[[Dict[str, Any]], bool]] = None,
-        **metadata_conditions
-    ):
-        """
-        检索相关文本
-        
-        Args:
-            raw_text: 查询文本
-            topk: 返回的最大结果数量
-            with_metadata: 是否包含元数据
-            index_name: 使用的索引名称
-            metadata_filter_func: 元数据过滤函数
-            **metadata_conditions: 元数据过滤条件
-        """
-        self.logger.debug(f"KVMemoryCollection: retrieve called")
-        
-        if index_name is None or index_name not in self.indexes:
-            self.logger.warning(f"Index '{index_name}' does not exist.")
-            return []
-        
-        if topk is None:
-            topk = self.default_topk
-        
-        index = self.indexes[index_name]["index"]
-        topk_ids = index.search(raw_text, topk=topk)
-        filtered_ids = self.filter_ids(topk_ids, metadata_filter_func, **metadata_conditions)
-        
-        self.logger.debug(f"检索到 {len(filtered_ids)} 条结果（请求 {topk} 条）")
-        
-        if with_metadata:
-            return [{"text": self.text_storage.get(i), "metadata": self.metadata_storage.get(i)}for i in filtered_ids]
-        else:
-            return [self.text_storage.get(i) for i in filtered_ids]
-    
-    def create_index(
-        self,
-        config: Optional[Dict[str, Any]] = None,
-        metadata_filter_func: Optional[Callable[[Dict[str, Any]], bool]] = None,
-        **metadata_conditions
-        ):
-        """
-        创建新索引
-        
-        Args:
-            config: 索引配置字典，必须包含name字段
-            metadata_filter_func: 元数据过滤函数
-            **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.indexes:
-            self.logger.warning(f"索引 '{index_name}' 已存在，无法重复创建")
-            return 0
-
-        # 从config中获取参数，如果没有则使用默认值
-        index_type = config.get("index_type", self.default_index_type)
-        description = config.get("description", f"Index for {index_name}")
-                
-        all_ids = self.get_all_ids()
-        filtered_ids = self.filter_ids(all_ids, metadata_filter_func, **metadata_conditions)
-        texts = [self.text_storage.get(i) for i in filtered_ids]
-            
-        try:
-            index = KVIndexFactory.create_index(
-                index_type=index_type,
-                index_name=index_name,
-                texts=texts,
-                ids=filtered_ids
-            )
-        except ValueError as e:
-            self.logger.error(f"Index type {index_type} not supported: {e}")
-            raise NotImplementedError(f"Index type {index_type} not supported: {e}")
-            
-        self.indexes[index_name] = {
-            "index": index,
-            "index_type": index_type,
-            "description": description,
-            "metadata_filter_func": metadata_filter_func,
-            "metadata_conditions": metadata_conditions,
-        }
-        
-        self.logger.info(f"成功创建索引 '{index_name}'，类型: {index_type}")
-        return 1  # 成功创建返回1
-        
-    def delete_index(self, index_name: str):
-        """
-        删除指定名称的索引
-        
-        Args:
-            index_name: 要删除的索引名称
-        """
-        if index_name in self.indexes:
-            del self.indexes[index_name]
-            self.logger.info(f"成功删除索引: {index_name}")
-        else:
-            raise ValueError(f"Index '{index_name}' does not exist.")
-    
-    def rebuild_index(self, index_name: str):
-        """
-        重建指定索引
-        
-        Args:
-            index_name: 要重建的索引名称
-        """
-        if index_name not in self.indexes:
-            self.logger.warning(f"Index '{index_name}' does not exist.")
-            return False
-        
-        info = self.indexes[index_name]
-        self.logger.info(f"开始重建索引: {index_name}")
-        
-        # 保存原始配置
-        original_config = {
-            "name": index_name,
-            "index_type": info["index_type"],
-            "description": info["description"]
-        }
-        
-        self.delete_index(index_name) 
-        self.create_index(
-            config=original_config,
-            metadata_filter_func=info["metadata_filter_func"],
-            **info["metadata_conditions"]
-        )
-        
-        self.logger.info(f"索引 '{index_name}' 重建完成")
-        return True
-    
-    def list_index(self) -> List[Dict[str, str]]:
-        """
-        列出当前所有索引及其描述信息。
-        返回结构：[{"name": ..., "description": ...}, ...]
-        """
-        return [
-            {"name": name, "description": info["description"]}
-            for name, info in self.indexes.items()
-        ]
-
-
-if __name__ == "__main__":
-    import tempfile
-    import time
-    
-    print("=== KVMemoryCollection 测试开始 ===")
-    
-    # 使用临时目录进行测试
-    with tempfile.TemporaryDirectory() as temp_dir:
-        print(f"测试目录: {temp_dir}")
-        
-        # 1. 创建集合（使用新的config方式）
-        print("\n1. 创建集合...")
-        config = {
-            "name": "test_collection",
-            "default_topk": 5,
-            "default_index_type": "bm25s"
-        }
-        collection = KVMemoryCollection(config)
-        print(f"✓ 集合创建成功: {collection.name}")
-        
-        # 2. 创建索引（使用新的config方式）
-        print("\n2. 创建索引...")
-        try:
-            index_config = {
-                "name": "main_index",
-                "index_type": "bm25s",
-                "description": "主要文本索引"
-            }
-            result = collection.create_index(config=index_config)
-            if result == 1:
-                print("✓ 主索引创建成功")
-            else:
-                print("✗ 主索引创建失败")
-        except Exception as e:
-            print(f"✗ 索引创建失败: {e}")
-        
-        # 3. 插入数据
-        print("\n3. 插入数据...")
-        test_data = [
-            ("这是第一条测试文本，包含人工智能相关内容。", {"category": "AI", "priority": 1}),
-            ("第二条文本讨论机器学习算法和深度学习。", {"category": "ML", "priority": 2}),
-            ("第三条关于自然语言处理技术的文档。", {"category": "NLP", "priority": 1}),
-            ("第四条涉及计算机视觉和图像识别。", {"category": "CV", "priority": 3}),
-            ("最后一条是关于数据科学的综合性文档。", {"category": "DS", "priority": 2})
-        ]
-        
-        inserted_ids = []
-        for text, metadata in test_data:
-            try:
-                doc_id = collection.insert(text, metadata, "main_index")
-                inserted_ids.append(doc_id)
-                print(f"✓ 插入成功: {text[:20]}...")
-            except Exception as e:
-                print(f"✗ 插入失败: {e}")
-        
-        print(f"✓ 总共插入了 {len(inserted_ids)} 条数据")
-        
-        # 4. 检索测试
-        print("\n4. 检索测试...")
-        search_queries = [
-            "人工智能", "机器学习", "自然语言处理", "计算机视觉", "数据科学"
-        ]
-        
-        for query in search_queries:
-            try:
-                results = collection.retrieve(
-                    raw_text=query,
-                    topk=3,
-                    with_metadata=True,
-                    index_name="main_index"
-                )
-                print(f"✓ '{query}' 找到 {len(results)} 条结果")
-                if results:
-                    top_result = results[0]["text"][:30] + "..." if len(results[0]["text"]) > 30 else results[0]["text"]
-                    print(f"  最相关: {top_result}")
-            except Exception as e:
-                print(f"✗ 检索 '{query}' 失败: {e}")
-        
-        # 5. 元数据过滤测试
-        print("\n5. 元数据过滤测试...")
-        try:
-            # 测试高优先级文档检索
-            results = collection.retrieve(
-                raw_text="技术",
-                topk=5,
-                with_metadata=True,
-                index_name="main_index",
-                priority=1  # 只检索优先级为1的文档
-            )
-            high_priority_count = len([r for r in results if r["metadata"].get("priority") == 1])
-            print(f"✓ 优先级过滤测试: 找到 {len(results)} 条结果，高优先级文档 {high_priority_count} 条")
-            
-        except Exception as e:
-            print(f"✗ 元数据过滤测试失败: {e}")
-        
-        # 6. 批量操作测试
-        print("\n6. 批量操作测试...")
-        try:
-            # 测试更新
-            if inserted_ids:
-                original_text = test_data[0][0]
-                new_text = "这是更新后的第一条文本，包含改进的人工智能内容。"
-                new_metadata = {"category": "AI", "priority": 1, "updated": True}
-                
-                new_id = collection.update(
-                    former_text=original_text,
-                    new_text=new_text,
-                    new_metadata=new_metadata
-                )
-                collection.insert(new_text, new_metadata, "main_index")
-                print(f"✓ 文档更新成功，新ID: {new_id[:8]}...")
-                
-                # 验证更新结果
-                results = collection.retrieve(
-                    raw_text="改进的人工智能",
-                    topk=1,
-                    with_metadata=True,
-                    index_name="main_index"
-                )
-                if results and results[0]["metadata"].get("updated"):
-                    print("✓ 更新验证成功")
-                else:
-                    print("✗ 更新验证失败")
-                
-                # 测试删除
-                collection.delete(new_text)
-                print("✓ 文档删除成功")
-                
-        except Exception as e:
-            print(f"✗ 批量操作测试失败: {e}")
-        
-        # 7. 索引管理测试
-        print("\n7. 索引管理测试...")
-        try:
-            # 创建分类索引（使用新的config方式）
-            ai_index_config = {
-                "name": "ai_index",
-                "index_type": "bm25s",
-                "description": "AI相关文档索引"
-            }
-            collection.create_index(
-                config=ai_index_config,
-                category="AI"  # 只包含AI类别的文档
-            )
-            
-            # 重新插入AI文档到新索引
-            for text, metadata in test_data:
-                if metadata.get("category") == "AI":
-                    collection.insert(text, metadata, "ai_index")
-            
-            # 列出索引
-            indexes = collection.list_index()
-            print(f"✓ 当前索引列表 ({len(indexes)} 个):")
-            for idx in indexes:
-                print(f"  - {idx['name']}: {idx['description']}")
-            
-            # 测试索引重建
-            collection.rebuild_index("main_index")
-            print("✓ 索引重建成功")
-            
-            # 测试索引删除
-            collection.delete_index("ai_index")
-            print("✓ 索引删除成功")
-            
-        except Exception as e:
-            print(f"✗ 索引管理测试失败: {e}")
-        
-        # 8. 持久化测试
-        print("\n8. 持久化测试...")
-        try:
-            # 保存集合
-            save_result = collection.store(temp_dir)
-            print(f"✓ 集合保存成功: {save_result['collection_path']}")
-            
-            # 加载集合
-            loaded_collection = KVMemoryCollection.load("test_collection", temp_dir)
-            print(f"✓ 集合加载成功: {loaded_collection.name}")
-            
-            # 验证加载的数据
-            if loaded_collection.indexes:
-                results = loaded_collection.retrieve(
-                    raw_text="机器学习",
-                    topk=2,
-                    index_name="main_index"
-                )
-                print(f"✓ 加载后检索测试成功，找到 {len(results)} 条结果")
-            else:
-                print("✗ 加载后索引为空")
-            
-        except Exception as e:
-            print(f"✗ 持久化测试失败: {e}")
-        
-        # 9. 错误处理测试
-        print("\n9. 错误处理测试...")
-        try:
-            # 测试不存在的索引
-            results = collection.retrieve(
-                raw_text="测试",
-                index_name="nonexistent_index",
-                topk=1
-            )
-            print(f"✓ 不存在索引处理: 返回 {len(results)} 条结果（预期为0）")
-            
-            # 测试删除不存在的索引
-            try:
-                collection.delete_index("nonexistent_index")
-                print("✗ 删除不存在索引应该失败")
-            except ValueError:
-                print("✓ 删除不存在索引正确抛出异常")
-            
-            # 测试重建不存在的索引
-            result = collection.rebuild_index("nonexistent_index")
-            if not result:
-                print("✓ 重建不存在索引正确返回False")
-            
-        except Exception as e:
-            print(f"✗ 错误处理测试失败: {e}")
-        
-        # 10. 清理测试
-        print("\n10. 清理测试...")
-        try:
-            KVMemoryCollection.clear("test_collection", temp_dir)
-            print("✓ 集合清理成功")
-            
-            # 验证清理结果
-            try:
-                KVMemoryCollection.load("test_collection", temp_dir)
-                print("✗ 清理后仍能加载集合")
-            except FileNotFoundError:
-                print("✓ 清理验证成功")
-        except Exception as e:
-            print(f"✗ 清理测试失败: {e}")
-    
-    print("\n=== 测试完成 ===")
-    print("✓ 所有主要功能测试通过")
-    print("注意: 某些测试可能由于依赖项（如KVIndexFactory）未完全加载而失败，这是正常的。")
-
-