isage-middleware 0.1.0__py3-none-any.whl → 0.1.4__py3-none-any.whl

sage/middleware/components/neuromem/memory_service.py

@@ -0,0 +1,324 @@
+"""
+Memory Service: 将Memory Manager包装成服务进程形式的封装层
+"""
+from typing import Dict, List, Any, Optional, Union
+from sage.middleware.components.neuromem.memory_manager import MemoryManager
+from sage.middleware.components.neuromem.memory_collection.base_collection import BaseMemoryCollection
+from sage.common.utils.logging.custom_logger import CustomLogger
+
+class MemoryService:
+    """
+    Memory服务封装，对外提供统一的服务接口
+    主要功能:
+    1. 管理collections的创建、删除、重命名等
+    2. 对collection进行数据操作(插入、查询、更新、删除)
+    3. 对collection进行索引操作(创建、删除、重建)
+    """
+    
+    def __init__(self, data_dir: Optional[str] = None):
+        self.manager = MemoryManager(data_dir)
+        self.logger = CustomLogger()
+        
+    def create_collection(self, name: str, backend_type: str, description: str = "", 
+                         embedding_model: Optional[Any] = None, dim: Optional[int] = None) -> Dict[str, Any]:
+        """创建新的collection"""
+        try:
+            collection = self.manager.create_collection(
+                name=name,
+                backend_type=backend_type,
+                description=description,
+                embedding_model=embedding_model,
+                dim=dim
+            )
+            # 预先注册常用的元数据字段
+            if hasattr(collection, "add_metadata_field"):
+                # 注册基础字段
+                collection.add_metadata_field("type")
+                collection.add_metadata_field("date")
+                collection.add_metadata_field("source")
+                collection.add_metadata_field("index")
+                collection.add_metadata_field("timestamp")
+                
+            return {
+                "status": "success",
+                "message": f"Collection '{name}' created successfully",
+                "collection_name": name
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def delete_collection(self, name: str) -> Dict[str, Any]:
+        """删除指定collection"""
+        try:
+            self.manager.delete_collection(name)
+            return {
+                "status": "success",
+                "message": f"Collection '{name}' deleted successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def list_collections(self) -> Dict[str, Any]:
+        """列出所有collections及其信息"""
+        try:
+            collections = []
+            for name, info in self.manager.collection_metadata.items():
+                collections.append({
+                    "name": name,
+                    "description": info.get("description", ""),
+                    "backend_type": info.get("backend_type", ""),
+                    "status": self.manager.collection_status.get(name, "unknown")
+                })
+            return {
+                "status": "success",
+                "collections": collections
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def rename_collection(self, old_name: str, new_name: str, 
+                         new_description: Optional[str] = None) -> Dict[str, Any]:
+        """重命名collection"""
+        try:
+            self.manager.rename(old_name, new_name, new_description)
+            return {
+                "status": "success",
+                "message": f"Collection renamed from '{old_name}' to '{new_name}'"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def get_collection_info(self, name: str) -> Dict[str, Any]:
+        """获取指定collection的详细信息"""
+        try:
+            info = self.manager.list_collection(name)
+            return {
+                "status": "success",
+                "collection_info": info
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    # Collection数据操作接口
+    def insert_data(self, collection_name: str, text: str, 
+                    metadata: Optional[Dict[str, Any]] = None,
+                    index_names: Optional[List[str]] = None) -> Dict[str, Any]:
+        """向指定collection插入数据"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            index_names = index_names or []
+            stable_id = collection.insert(text, metadata, *index_names)
+            
+            return {
+                "status": "success",
+                "message": "Data inserted successfully",
+                "id": stable_id
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def update_data(self, collection_name: str, old_text: str, new_text: str,
+                    new_metadata: Optional[Dict[str, Any]] = None,
+                    index_names: Optional[List[str]] = None) -> Dict[str, Any]:
+        """更新指定collection中的数据"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            index_names = index_names or []
+            new_id = collection.update(old_text, new_text, new_metadata, *index_names)
+            
+            return {
+                "status": "success",
+                "message": "Data updated successfully",
+                "new_id": new_id
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def delete_data(self, collection_name: str, text: str) -> Dict[str, Any]:
+        """从指定collection删除数据"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            collection.delete(text)
+            return {
+                "status": "success",
+                "message": "Data deleted successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def retrieve_data(self, collection_name: str, query_text: str,
+                      topk: Optional[int] = None,
+                      index_name: Optional[str] = None,
+                      with_metadata: bool = False,
+                      metadata_filter: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        """从指定collection检索数据"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+            
+            metadata_conditions = metadata_filter or {}
+            results = collection.retrieve(
+                query_text,
+                topk=topk,
+                index_name=index_name,
+                with_metadata=with_metadata,
+                **metadata_conditions
+            )
+            
+            return {
+                "status": "success",
+                "results": results
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    # 索引操作接口
+    def create_index(self, collection_name: str, index_name: str,
+                     description: Optional[str] = None,
+                     metadata_conditions: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        """为指定collection创建索引"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+            
+            metadata_conditions = metadata_conditions or {}
+            collection.create_index(
+                index_name=index_name,
+                description=description,
+                **metadata_conditions
+            )
+            
+            return {
+                "status": "success",
+                "message": f"Index '{index_name}' created successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def delete_index(self, collection_name: str, index_name: str) -> Dict[str, Any]:
+        """删除指定collection的索引"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            collection.delete_index(index_name)
+            return {
+                "status": "success",
+                "message": f"Index '{index_name}' deleted successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def rebuild_index(self, collection_name: str, index_name: str) -> Dict[str, Any]:
+        """重建指定collection的索引"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            collection.rebuild_index(index_name)
+            return {
+                "status": "success",
+                "message": f"Index '{index_name}' rebuilt successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def list_indexes(self, collection_name: str) -> Dict[str, Any]:
+        """列出指定collection的所有索引"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            indexes = collection.list_index()
+            return {
+                "status": "success",
+                "indexes": indexes
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    # 存储操作接口
+    def store_collection(self, collection_name: str) -> Dict[str, Any]:
+        """保存指定collection到磁盘"""
+        try:
+            collection = self.manager.get_collection(collection_name)
+            if not collection:
+                raise ValueError(f"Collection '{collection_name}' not found")
+                
+            self.manager.store_collection(collection_name)
+            return {
+                "status": "success",
+                "message": f"Collection '{collection_name}' stored successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }
+
+    def store(self) -> Dict[str, Any]:
+        """保存整个manager的所有信息到磁盘"""
+        try:
+            self.manager.store_collection()  # 保存所有已加载的collection
+            return {
+                "status": "success",
+                "message": "All manager data stored successfully"
+            }
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": str(e)
+            }

sage/middleware/components/neuromem/micro_service/__init__.py

@@ -0,0 +1,37 @@
+"""
+SAGE - Streaming-Augmented Generative Execution
+"""
+
+# 动态版本加载
+def _load_version():
+    """从项目根目录动态加载版本信息"""
+    from pathlib import Path
+    
+    # 获取项目根目录
+    current_file = Path(__file__).resolve()
+    root_dir = current_file.parent.parent.parent.parent.parent.parent.parent.parent
+    version_file = root_dir / "_version.py"
+    
+    # 加载版本信息
+    if version_file.exists():
+        version_globals = {}
+        with open(version_file, 'r', encoding='utf-8') as f:
+            exec(f.read(), version_globals)
+        return {
+            'version': version_globals.get('__version__', '0.1.4'),
+            'author': version_globals.get('__author__', 'SAGE Team'),
+            'email': version_globals.get('__email__', 'shuhao_zhang@hust.edu.cn')
+        }
+    
+    # 默认值
+    return {
+        'version': '0.1.4',
+        'author': 'SAGE Team', 
+        'email': 'shuhao_zhang@hust.edu.cn'
+    }
+
+# 加载信息
+_info = _load_version()
+__version__ = _info['version']
+__author__ = _info['author']
+__email__ = _info['email']

sage/middleware/components/neuromem/micro_service/neuromem_vdb.py

@@ -0,0 +1,198 @@
+from logging import Manager
+import os
+from re import M
+from typing import Dict, List, Any, Optional, Union
+from sage.middleware.components.neuromem.memory_manager import MemoryManager
+from sage.middleware.components.neuromem.memory_collection.vdb_collection import VDBMemoryCollection
+
+class NeuroMemVDB:
+    def __init__(self):
+        self.manager = MemoryManager(self._get_default_data_dir())
+        self.online_register_collection: Dict[str, VDBMemoryCollection] = {}
+
+
+    def register_collection(self, collection_name: str, config = None):
+        # 检查collection是否已存在
+        existing_collections = self.manager.list_collection()
+        collection_exists = any(c['name'] == collection_name for c in existing_collections) if isinstance(existing_collections, list) else False
+        
+        if collection_exists:
+            if config is not None:
+                print(f"警告: Collection '{collection_name}' 已存在，忽略传入的config参数")
+            # 正常连接并注册到online_register_collection
+            collection = self.manager.get_collection(collection_name)
+            if collection:
+                self.online_register_collection[collection_name] = collection
+                print(f"成功连接已存在的collection: {collection_name}")
+        else:
+            if config is not None:
+                # 正常通过Manager创建并注册到online_register_collection
+                embedding_model = config.get('embedding_model')
+                dim = config.get('dim')
+                description = config.get('description', f"VDB collection: {collection_name}")
+                collection = self.manager.create_collection(
+                    name=collection_name,
+                    backend_type="VDB",
+                    description=description,
+                    embedding_model=embedding_model,
+                    dim=dim
+                )
+                if collection:
+                    self.online_register_collection[collection_name] = collection
+                    print(f"成功创建新collection: {collection_name}")
+            else:
+                # 警告 用默认参数创建一个collection并注册到online_register_collection
+                print(f"警告: Collection '{collection_name}' 不存在且未提供config，使用默认参数创建")
+                collection = self.manager.create_collection(
+                    name=collection_name,
+                    backend_type="VDB",
+                    description=f"Default VDB collection: {collection_name}"
+                )
+                if collection:
+                    self.online_register_collection[collection_name] = collection
+                    print(f"成功创建默认collection: {collection_name}")
+
+    def insert(self, raw_data: Any, metadata: Optional[Dict[str, Any]] = None):
+        # 该方法在所有 online_register_collection 均插入数据
+        if not self.online_register_collection:
+            print("警告: 没有注册的collection，无法插入数据")
+            return
+        
+        results = {}
+        for collection_name, collection in self.online_register_collection.items():
+            try:
+                stable_id = collection.insert(raw_data, metadata)
+                results[collection_name] = stable_id
+                print(f"成功在collection '{collection_name}' 中插入数据，ID: {stable_id}")
+            except Exception as e:
+                print(f"在collection '{collection_name}' 插入数据失败: {str(e)}")
+                results[collection_name] = None
+        return results
+    
+    def retrieve(self, raw_data: Any, topk: int = 5):
+        # 测试类型的retrieve方法，不返回任何值，直接print出检索结果
+        if not self.online_register_collection:
+            print("警告: 没有注册的collection，无法检索数据")
+            return
+        
+        print(f"检索查询: {raw_data}")
+        print("=" * 50)
+        
+        for collection_name, collection in self.online_register_collection.items():
+            try:
+                results = collection.retrieve(raw_data, topk=topk, with_metadata=True)
+                print(f"Collection '{collection_name}' 检索结果:")
+                if results:
+                    for i, result in enumerate(results, 1):
+                        print(f"  {i}. 文本: {result['text']}")
+                        if result['metadata']:
+                            print(f"     元数据: {result['metadata']}")
+                else:
+                    print("  无结果")
+                print("-" * 30)
+            except Exception as e:
+                print(f"Collection '{collection_name}' 检索失败: {str(e)}")
+                print("-" * 30)
+
+    def build_index(self, index_name: str = "global_index", description: str = "全局索引"):
+        # 该方法在所有 online_register_collection 创建指定索引
+        if not self.online_register_collection:
+            print("警告: 没有注册的collection，无法创建索引")
+            return
+        
+        results = {}
+        for collection_name, collection in self.online_register_collection.items():
+            try:
+                collection.create_index(index_name, description=description)
+                results[collection_name] = True
+                print(f"成功在collection '{collection_name}' 中创建索引 '{index_name}'")
+            except Exception as e:
+                print(f"在collection '{collection_name}' 创建索引失败: {str(e)}")
+                results[collection_name] = False
+        return results
+
+    def store_to_disk(self):
+        # 调用manager的store_collection()方法，触发一次全局落盘
+        try:
+            self.manager.store_collection()
+            print("成功将所有collection数据保存到磁盘")
+        except Exception as e:
+            print(f"保存数据到磁盘失败: {str(e)}")
+
+    def clean_collection(self, collection_name: str):
+        # 清空某个collection的所有数据(包括磁盘中存在的)
+        try:
+            # 从在线注册列表中移除
+            if collection_name in self.online_register_collection:
+                del self.online_register_collection[collection_name]
+                print(f"已从在线注册列表中移除collection: {collection_name}")
+            
+            # 通过manager删除collection（包括磁盘数据）
+            self.manager.delete_collection(collection_name)
+            print(f"成功清理collection: {collection_name}")
+        except Exception as e:
+            print(f"清理collection '{collection_name}' 失败: {str(e)}")
+
+    @classmethod
+    def _get_default_data_dir(cls):
+        """获取默认数据目录"""
+        cur_dir = os.getcwd()
+        data_dir = os.path.join(cur_dir, "data", "neuromem_vdb")
+        os.makedirs(data_dir, exist_ok=True)
+        return data_dir
+    
+if __name__ == "__main__":
+    print("默认数据目录:", NeuroMemVDB._get_default_data_dir())
+    print("\n" + "="*60)
+    print("开始测试 NeuroMemVDB")
+    print("="*60)
+    
+    # 创建NeuroMemVDB实例
+    vdb = NeuroMemVDB()
+    
+    # 测试1: 注册不存在的collection（无config，应该警告并使用默认参数）
+    print("\n1. 测试注册不存在的collection（无config）")
+    vdb.register_collection("test_collection1")
+    
+    # 测试2: 注册不存在的collection（有config）
+    print("\n2. 测试注册不存在的collection（有config）")
+    config = {
+        'description': '测试用的VDB集合',
+        'embedding_model': None,  # 使用默认embedding模型
+        'dim': None  # 使用默认维度
+    }
+    vdb.register_collection("test_collection2", config)
+    
+    # 测试3: 插入数据
+    print("\n3. 测试插入数据")
+    vdb.insert("Python是一种编程语言", {"type": "test", "priority": "high"})
+    vdb.insert("机器学习是人工智能的重要分支", {"type": "demo", "priority": "low"})
+    vdb.insert("向量数据库用于存储和检索高维向量数据")
+    
+    # 测试4: 创建索引
+    print("\n4. 测试创建索引")
+    vdb.build_index("custom_index", "自定义测试索引")
+    
+    # 测试5: 检索数据
+    print("\n5. 测试检索数据")
+    vdb.retrieve("编程语言", topk=3)
+    
+    # 测试6: 保存到磁盘
+    print("\n6. 测试保存到磁盘")
+    vdb.store_to_disk()
+    
+    # 测试7: 重新注册已存在的collection（应该警告）
+    print("\n7. 测试重新注册已存在的collection")
+    vdb.register_collection("test_collection1", {"description": "这个config应该被忽略"})
+    
+    # 测试8: 清理collection
+    print("\n8. 测试清理collection")
+    vdb.clean_collection("test_collection1")
+    
+    # 测试9: 再次检索（应该只有test_collection2的结果）
+    print("\n9. 清理后再次检索")
+    vdb.retrieve("编程语言", topk=3)
+    
+    print("\n" + "="*60)
+    print("测试完成")
+    print("="*60)