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

sage/middleware/components/neuromem/memory_service.py

@@ -1,324 +0,0 @@
-"""
-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

@@ -1,56 +0,0 @@
-"""
-SAGE - Streaming-Augmented Generative Execution
-"""
-
-# 动态版本加载
-def _load_version():
-    """从 sage-common 包加载版本信息"""
-    try:
-        # 优先从 sage-common 包加载版本
-        from sage.common._version import __version__
-        return {
-            'version': __version__,
-            'author': 'SAGE Team',
-            'email': 'shuhao_zhang@hust.edu.cn'
-        }
-    except ImportError:
-        # 如果 sage-common 不可用，从项目根目录加载（开发环境）
-        try:
-            from pathlib import Path
-            current_file = Path(__file__).resolve()
-            # 根据当前文件位置计算到项目根目录的层数
-            parts = current_file.parts
-            sage_index = -1
-            for i, part in enumerate(parts):
-                if part == 'SAGE':
-                    sage_index = i
-                    break
-            
-            if sage_index >= 0:
-                root_dir = Path(*parts[:sage_index+1])
-                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.3'),
-                        'author': version_globals.get('__author__', 'SAGE Team'),
-                        'email': version_globals.get('__email__', 'shuhao_zhang@hust.edu.cn')
-                    }
-        except Exception:
-            pass
-    
-    # 最后的默认值
-    return {
-        'version': '0.1.3',
-        '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

@@ -1,198 +0,0 @@
-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)