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

sage/middleware/components/neuromem/memory_manager.py

@@ -0,0 +1,401 @@
+# sage/service/memory/memory_manager.py
+# python -m sage.middleware.services.neuromem.memory_manager
+import os
+import json
+import shutil
+from typing import Any, Dict, List, Optional, Union
+
+from sage.common.utils.logging.custom_logger import CustomLogger
+from sage.middleware.components.neuromem.utils.path_utils import get_default_data_dir
+from sage.middleware.components.neuromem.memory_collection.base_collection import BaseMemoryCollection
+from sage.middleware.components.neuromem.memory_collection.graph_collection import GraphMemoryCollection
+from sage.middleware.components.neuromem.memory_collection.kv_collection import KVMemoryCollection
+from sage.middleware.components.neuromem.memory_collection.vdb_collection import VDBMemoryCollection
+from sage.middleware.utils.embedding.embedding_api import apply_embedding_model
+
+class MemoryManager:
+    """
+    内存管理器，管理不同类型的 MemoryCollection 实例
+    """
+
+    def __init__(self, data_dir: Optional[str] = None):
+        self.logger = CustomLogger()
+        
+        self.collections: Dict[str, BaseMemoryCollection] = {}
+        self.collection_metadata: Dict[str, Dict[str, Any]] = {}
+        self.collection_status: Dict[str, str] = {}  # 状态表: "loaded" or "on_disk"
+        
+        # 从默认目录data/sage_memory获取disk中的manager(如果有)
+        self.data_dir = data_dir or get_default_data_dir()
+        self.manager_path = os.path.join(self.data_dir, "manager.json")
+        self._load_manager()  
+
+    def create_collection(
+            self,
+            name: str,
+            backend_type: str,
+            description: str = "",
+            embedding_model: Optional[Any] = None,
+            dim: Optional[int] = None,
+    ) -> BaseMemoryCollection:
+        """
+        创建新的集合
+        直接返回collection
+        """
+        
+        if name in self.collection_metadata:
+            self.logger.warning(f"Collection with name '{name}' already exists.")
+            return None
+
+        metadata = {
+        "description": description,
+        "backend_type": backend_type
+        }
+
+        # 创建基础集合
+        if backend_type == "VDB":
+            if embedding_model is None or dim is None:
+                self.logger.warning(f"VDB requires parameter 'embedding_model' and 'dim', use default values.")
+                embedding_model = apply_embedding_model("default")
+                dim = embedding_model.get_dim()
+
+            # 使用新的config字典方式创建VDB集合
+            vdb_config = {
+                "name": name,
+                "default_embedding_model": embedding_model.method_name,
+                "default_dim": dim
+            }
+            new_collection = VDBMemoryCollection(vdb_config)
+            metadata.update({
+                "embedding_model_method": embedding_model.method_name,
+                "embedding_dim": dim
+            })
+        elif backend_type == "KV":
+            # 使用新的config字典方式创建KV集合
+            kv_config = {
+                "name": name
+            }
+            new_collection = KVMemoryCollection(kv_config)
+        elif backend_type == "GRAPH":
+            new_collection = GraphMemoryCollection(name)
+        else:
+            self.logger.error(f"Unsupported backend_type: {backend_type}")
+
+        # 存储到 collections
+        self.collections[name] = new_collection
+        self.collection_metadata[name] = metadata
+        self.collection_status[name] = "loaded"
+        return new_collection
+
+    def get_collection(self, name: str) -> Optional[BaseMemoryCollection]:
+        """优先返回内存collection，不在内存则尝试磁盘懒加载并发警告"""
+        if name in self.collections:
+            return self.collections[name]
+        elif name in self.collection_metadata:
+            self.logger.warning(f"Collection '{name}' not in memory, loading from disk.")
+            return self._load_collection(name)
+        else:
+            self.logger.warning(f"Collection '{name}' not found.")
+            return None
+
+    def delete_collection(self, name: str):
+        """删除collection：内存、元数据、磁盘三处都要删"""
+        existed = False
+        # 内存
+        if name in self.collections:
+            del self.collections[name]
+            existed = True
+        # 元信息
+        if name in self.collection_metadata:
+            backend_type = self.collection_metadata[name].get("backend_type")
+            del self.collection_metadata[name]
+            self.collection_status.pop(name, None)
+            existed = True
+            # 磁盘
+            if backend_type == "VDB":
+                path = os.path.join(self.data_dir, "vdb_collection", name)
+            elif backend_type == "KV":
+                path = os.path.join(self.data_dir, "kv_collection", name)
+            elif backend_type == "GRAPH":
+                path = os.path.join(self.data_dir, "graph_collection", name)
+            else:
+                path = None
+            if path and os.path.exists(path):
+                try:
+                    if os.path.isdir(path):
+                        shutil.rmtree(path)
+                    else:
+                        os.remove(path)
+                except Exception as e:
+                    self.logger.warning(f"Failed to delete disk collection '{name}': {e}")
+        if not existed:
+            self.logger.warning(f"Collection '{name}' not found.")
+        self._save_manager()
+        
+    def connect_collection(self, name: str, embedding_model=None) -> Optional[BaseMemoryCollection]:
+        """类似get_collection，但支持外部embedding_model"""
+        if name in self.collections:
+            return self.collections[name]
+        if name not in self.collection_metadata:
+            self.logger.warning(f"Collection '{name}' not found in metadata.")
+            return None
+        meta = self.collection_metadata[name]
+        backend_type = meta.get("backend_type")
+        if backend_type == "VDB":
+            vdb_path = os.path.join(self.data_dir, "vdb_collection", name)
+            if embedding_model is None:
+                method_name = meta.get("embedding_model_method", "default")
+                embedding_model = apply_embedding_model(method_name)
+            # 使用新的load方法，只传入name和路径
+            collection = VDBMemoryCollection.load(name, vdb_path)
+        elif backend_type == "KV":
+            kv_path = os.path.join(self.data_dir, "kv_collection", name)
+            # 使用新的load方法，只传入name和路径
+            collection = KVMemoryCollection.load(name, kv_path)
+        elif backend_type == "GRAPH":
+            graph_path = os.path.join(self.data_dir, "graph_collection", name)
+            collection = GraphMemoryCollection.load(name, graph_path)
+        else:
+            self.logger.warning(f"Unknown backend_type: {backend_type}")
+            return None
+        self.collections[name] = collection
+        self.collection_status[name] = "loaded"
+        return collection
+
+    def store_collection(self, name: Optional[str] = None):
+        """
+        只保存已加载（即被用过的）的collection及manager信息
+        """
+        # 指定name只存一个，否则只存所有已加载collection
+        collections_to_save = [name] if name else [
+            n for n, status in self.collection_status.items() if status == "loaded"
+        ]
+        for cname in collections_to_save:
+            if cname not in self.collections:
+                continue
+            collection = self.collections[cname]
+            if hasattr(collection, "store"):
+                # 传递MemoryManager的数据目录，让collection使用统一的目录结构
+                collection.store(self.data_dir)
+        self._save_manager()
+        
+    def _save_manager(self):
+        with open(self.manager_path, "w", encoding="utf-8") as f:
+            json.dump(self.collection_metadata, f, ensure_ascii=False, indent=2)
+        self.logger.info(f"Manager info saved to {self.manager_path}")  
+          
+    def _load_manager(self):
+        """只加载collection_metadata和状态，不加载任何collection本体"""
+        if not os.path.exists(self.manager_path):
+            return
+        with open(self.manager_path, "r", encoding="utf-8") as f:
+            self.collection_metadata = json.load(f)
+        for name in self.collection_metadata:
+            self.collection_status[name] = "on_disk"
+    
+    def _load_collection(self, name: str) -> Optional[BaseMemoryCollection]:
+        """懒加载collection进内存"""
+        if name not in self.collection_metadata:
+            self.logger.warning(f"Collection '{name}' metadata not found in manager.")
+            return None
+
+        meta = self.collection_metadata[name]
+        backend_type = meta.get("backend_type")
+        if backend_type == "VDB":
+            vdb_path = os.path.join(self.data_dir, "vdb_collection", name)
+            method_name = meta.get("embedding_model_method", "default")
+            embedding_model = apply_embedding_model(method_name)
+            # 使用新的load方法，只传入name和路径
+            collection = VDBMemoryCollection.load(name, vdb_path)
+        elif backend_type == "KV":
+            kv_path = os.path.join(self.data_dir, "kv_collection", name)
+            # 使用新的load方法，只传入name和路径
+            collection = KVMemoryCollection.load(name, kv_path)
+        elif backend_type == "GRAPH":
+            graph_path = os.path.join(self.data_dir, "graph_collection", name)
+            collection = GraphMemoryCollection.load(name, graph_path)
+        else:
+            self.logger.warning(f"Unknown backend_type: {backend_type}")
+            return None
+
+        self.collections[name] = collection
+        self.collection_status[name] = "loaded"
+        return collection
+    
+            
+    def list_collection(self, name: Optional[str] = None) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+        """
+        列出一个或所有 collection 的基本信息。
+        List basic info of one or all collections.
+        """
+        if name:
+            if name not in self.collection_metadata:
+                self.logger.warning(f"Collection '{name}' not found.")
+                return None
+            return {"name": name, "status": self.collection_status.get(name, "unknown"), **self.collection_metadata[name]}
+        else:
+            return [
+                {"name": n, "status": self.collection_status.get(n, "unknown"), **meta}
+                for n, meta in self.collection_metadata.items()
+            ]
+
+    def rename(self, former_name: str, new_name: str, new_description: Optional[str] = None):
+        """重命名 collection 并更新描述"""
+        if former_name not in self.collection_metadata:
+            self.logger.warning(f"Collection '{former_name}' not found.")
+            return False
+        if new_name in self.collection_metadata:
+            self.logger.warning(f"Collection '{new_name}' already exists.")
+            return False
+
+        # 重命名内存对象
+        if former_name in self.collections:
+            collection = self.collections.pop(former_name)
+            collection.name = new_name
+            self.collections[new_name] = collection
+
+        # 重命名状态表
+        if former_name in self.collection_status:
+            self.collection_status[new_name] = self.collection_status.pop(former_name)
+
+        # 更新元信息
+        metadata = self.collection_metadata.pop(former_name)
+        metadata["description"] = new_description or metadata.get("description", "")
+        self.collection_metadata[new_name] = metadata
+
+        # 重命名磁盘目录/文件
+        backend_type = metadata.get("backend_type")
+        def get_path(n):
+            if backend_type == "VDB":
+                return os.path.join(self.data_dir, "vdb_collection", n)
+            elif backend_type == "KV":
+                return os.path.join(self.data_dir, "kv_collection", n)
+            elif backend_type == "GRAPH":
+                return os.path.join(self.data_dir, "graph_collection", n)
+        old_path = get_path(former_name)
+        new_path = get_path(new_name)
+        if os.path.exists(old_path):
+            os.rename(old_path, new_path)
+        self._save_manager()
+        return True
+
+
+
+if __name__ == "__main__":
+    import os
+    import shutil
+
+    def print_result(expect, actual, passed):
+        green = '\033[92m'
+        red = '\033[91m'
+        endc = '\033[0m'
+        print(f"预期结果：{expect}")
+        print(f"实际结果：{actual}")
+        print("是否通过测试：", end="")
+        print(f"{green}是{endc}" if passed else f"{red}否{endc}")
+        print("="*50)
+
+    def clear_all_data(manager):
+        # 删除data_dir下所有collection目录和manager.json
+        base = manager.data_dir
+        for subdir in ["vdb_collection"]:
+            path = os.path.join(base, subdir)
+            if os.path.exists(path):
+                shutil.rmtree(path)
+        if os.path.exists(manager.manager_path):
+            os.remove(manager.manager_path)
+
+    # 测试1: 创建新VDB collection
+    manager = MemoryManager()
+    c1 = manager.create_collection(name="test_vdb", backend_type="VDB")
+    passed = c1 is not None and "test_vdb" in manager.collections
+    print_result("新建test_vdb后能在collections中查到", str("test_vdb" in manager.collections), passed)
+
+    # 测试2: 再次新建同名collection应警告并返回None
+    c2 = manager.create_collection(name="test_vdb", backend_type="VDB")
+    passed = c2 is None
+    print_result("新建已存在collection应返回None", str(c2 is None), passed)
+
+    # 测试3: list_collection返回正确，状态为loaded
+    info = manager.list_collection("test_vdb")
+    passed = info is not None and info['status'] == "loaded"
+    print_result("list_collection显示状态为loaded", f"{info}", passed)
+
+    # 测试4: get_collection取已存在的
+    c3 = manager.get_collection("test_vdb")
+    passed = c3 is c1
+    print_result("get_collection能返回内存已加载对象", f"id相同: {c3 is c1}", passed)
+
+    # 测试5: store_collection只保存已加载collection
+    manager.store_collection()
+    path = os.path.join(manager.data_dir, "vdb_collection", "test_vdb")
+    passed = os.path.exists(path)
+    print_result("store_collection后磁盘应有数据目录", str(os.path.exists(path)), passed)
+
+    # 测试6: 模拟“只在磁盘不在内存”的懒加载
+    del manager.collections["test_vdb"]
+    manager.collection_status["test_vdb"] = "on_disk"
+    passed = "test_vdb" not in manager.collections and os.path.exists(os.path.join(manager.data_dir, "vdb_collection", "test_vdb"))
+    print_result("未加载collection不会在collections", str(passed), passed)
+
+    # 测试7: get_collection能从磁盘懒加载
+    c4 = manager.get_collection("test_vdb")
+    passed = c4 is not None and "test_vdb" in manager.collections and manager.collection_status["test_vdb"] == "loaded"
+    print_result("get_collection能懒加载on_disk collection", str(passed), passed)
+
+    # 测试8: rename
+    r = manager.rename("test_vdb", "renamed_vdb", "renamed desc")
+    passed = r and "renamed_vdb" in manager.collections
+    print_result("rename操作成功", str(r), passed)
+
+    # store 一下，保证磁盘有数据
+    manager.store_collection("renamed_vdb")
+
+    print("\n==中间测试==")
+    print("请确认manager.json、collection磁盘文件已经存在。")
+    print("输入yes继续从磁盘加载测试，否则中止（将自动清理）:")
+    ans = input().strip().lower()
+    if ans != "yes":
+        print("用户中止，已清理所有测试数据。")
+        clear_all_data(manager)
+        exit(0)
+
+    # 测试9: 用新manager对象从磁盘懒加载
+    manager2 = MemoryManager(manager.data_dir)
+    loaded = "renamed_vdb" in manager2.collection_metadata and manager2.collection_status.get("renamed_vdb") == "on_disk"
+    print_result("新manager能从磁盘加载管理信息（collection未自动加载）", str(loaded), loaded)
+
+    # 测试10: list_collection（全部）
+    all_info = manager.list_collection()
+    passed = isinstance(all_info, list)
+    print_result("list_collection()能返回所有collection列表", str(passed), passed)
+
+    print("\n==中间测试==")
+    print("请确认manager.json、collection磁盘文件已经存在。")
+    print("输入yes继续从磁盘加载测试，否则中止（将自动清理）:")
+    ans = input().strip().lower()
+    if ans != "yes":
+        print("用户中止，已清理所有测试数据。")
+        clear_all_data(manager)
+        exit(0)
+
+    # 测试11: 用新manager对象从磁盘懒加载
+    manager2 = MemoryManager(manager.data_dir)
+    loaded = "test_vdb" in manager2.collection_metadata or "renamed_vdb" in manager2.collection_metadata
+    print_result("新manager能从磁盘加载管理信息（collection未自动加载）", str(loaded), loaded)
+
+    # 如果上面已重命名并删除，理论上只有空的或剩余collection
+    # 为了测试懒加载，再新建并保存一次
+    cnew = manager2.create_collection("lazy_test", backend_type="VDB")
+    manager2.store_collection()
+    del manager2.collections["lazy_test"]
+    manager2.collection_status["lazy_test"] = "on_disk"
+    c5 = manager2.get_collection("lazy_test")
+    passed = c5 is not None and manager2.collection_status["lazy_test"] == "loaded"
+    print_result("新manager懒加载on_disk collection后状态变为loaded", str(passed), passed)
+
+    # 清理所有
+    clear_all_data(manager)
+    print("\n所有测试完成，已清理测试数据！")
+
+