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

sage/middleware/components/neuromem/tests/test_memory_service.py

@@ -0,0 +1,293 @@
+"""
+测试Memory Service的集成测试用例
+"""
+import os
+import shutil
+import traceback
+from sage.middleware.components.neuromem.memory_service import MemoryService
+from sage.middleware.utils.embedding.embedding_api import apply_embedding_model
+
+def test_memory_service():
+    """测试Memory Service的主要功能（直接测试，不使用服务框架）"""
+    print("🚀 Starting Memory Service test...")
+    
+    try:
+        # 1. 直接创建MemoryService实例
+        # 使用默认的embedding model
+        embedding_model = apply_embedding_model("default")
+        dim = embedding_model.get_dim()
+        # 指定临时测试目录
+        test_data_dir = os.path.join(os.path.dirname(__file__), "test_data")
+        os.makedirs(test_data_dir, exist_ok=True)
+        memory_service = MemoryService(data_dir=test_data_dir)
+        
+        print("✅ Memory service created, testing operations...")
+        
+        # 2. 测试创建collection
+        result1 = memory_service.create_collection(
+            name="test_collection",
+            backend_type="VDB",
+            description="Test collection",
+            embedding_model=embedding_model,
+            dim=dim
+        )
+        print(f"Create collection result: {result1}")
+        assert result1["status"] == "success", f"Create collection failed: {result1}"
+        
+        # 3. 测试插入数据
+        result2 = memory_service.insert_data(
+            collection_name="test_collection",
+            text="This is a test document",
+            metadata={"type": "test", "date": "2025-07-26"}
+        )
+        print(f"Insert data result: {result2}")
+        assert result2["status"] == "success", f"Insert data failed: {result2}"
+        
+        # 4. 测试创建索引
+        result3 = memory_service.create_index(
+            collection_name="test_collection",
+            index_name="test_index",
+            description="Test index"
+        )
+        print(f"Create index result: {result3}")
+        assert result3["status"] == "success", f"Create index failed: {result3}"
+        
+        # 5. 测试检索数据
+        result4 = memory_service.retrieve_data(
+            collection_name="test_collection",
+            query_text="test document",
+            topk=5,
+            index_name="test_index",
+            with_metadata=True
+        )
+        print(f"Retrieve data result: {result4}")
+        assert result4["status"] == "success", f"Retrieve data failed: {result4}"
+        
+        # 6. 测试插入更多数据
+        for i in range(3):
+            result = memory_service.insert_data(
+                collection_name="test_collection",
+                text=f"Test document {i}",
+                metadata={"type": "test", "index": i}
+            )
+            print(f"Insert data {i} result: {result}")
+            assert result["status"] == "success", f"Insert data {i} failed: {result}"
+        
+        # 7. 测试列出collections
+        final_result = memory_service.list_collections()
+        print(f"Collections list: {final_result}")
+        assert final_result["status"] == "success", f"List collections failed: {final_result}"
+        assert len(final_result["collections"]) == 1, "Should have exactly 1 collection"
+        
+        # 8. 测试获取collection信息
+        info_result = memory_service.get_collection_info("test_collection")
+        print(f"Collection info: {info_result}")
+        assert info_result["status"] == "success", f"Get collection info failed: {info_result}"
+        
+        # 9. 测试列出索引
+        index_result = memory_service.list_indexes("test_collection")
+        print(f"Indexes list: {index_result}")
+        assert index_result["status"] == "success", f"List indexes failed: {index_result}"
+        
+        # 10. 测试存储单个collection
+        store_collection_result = memory_service.store_collection("test_collection")
+        print(f"Store collection result: {store_collection_result}")
+        assert store_collection_result["status"] == "success", f"Store collection failed: {store_collection_result}"
+        
+        # 11. 创建第二个collection来测试选择性存储
+        result_col2 = memory_service.create_collection(
+            name="temp_collection",
+            backend_type="VDB",
+            description="Temporary test collection",
+            embedding_model=embedding_model,
+            dim=dim
+        )
+        print(f"Create temp collection result: {result_col2}")
+        assert result_col2["status"] == "success", f"Create temp collection failed: {result_col2}"
+        
+        # 12. 向第二个collection插入数据
+        temp_insert_result = memory_service.insert_data(
+            collection_name="temp_collection",
+            text="This is temporary data",
+            metadata={"type": "temp", "temporary": True}
+        )
+        print(f"Insert temp data result: {temp_insert_result}")
+        assert temp_insert_result["status"] == "success", f"Insert temp data failed: {temp_insert_result}"
+        
+        # 13. 测试存储所有manager数据
+        store_all_result = memory_service.store()
+        print(f"Store all result: {store_all_result}")
+        assert store_all_result["status"] == "success", f"Store all failed: {store_all_result}"
+        
+        print("✅ All operations including storage completed successfully!")
+        
+        return True
+        
+    except Exception as e:
+        print(f"❌ Test failed: {e}")
+        traceback.print_exc()
+        return False
+        
+    finally:
+        # 清理资源
+        try:
+            # 清理测试数据目录
+            test_data_dir = os.path.join(os.path.dirname(__file__), "test_data")
+            if os.path.exists(test_data_dir):
+                shutil.rmtree(test_data_dir)
+            print("🧹 Cleanup completed")
+        except Exception as e:
+            print(f"⚠️ Cleanup error: {e}")
+
+
+def test_data_loss_without_storage():
+    """测试未保存collection导致的数据丢失"""
+    print("\n🔍 Starting data loss test without storage...")
+    
+    try:
+        # 使用默认的embedding model
+        embedding_model = apply_embedding_model("default")
+        dim = embedding_model.get_dim()
+        
+        # 创建独立的测试目录
+        test_data_dir = os.path.join(os.path.dirname(__file__), "test_data_loss")
+        os.makedirs(test_data_dir, exist_ok=True)
+        
+        print("✅ Starting data loss test...")
+        
+        # 阶段1: 创建两个collection并插入数据
+        print("\n📝 Phase 1: Creating collections and inserting data...")
+        
+        memory_service = MemoryService(data_dir=test_data_dir)
+        
+        # 创建第一个collection
+        result1 = memory_service.create_collection(
+            name="saved_collection",
+            backend_type="VDB", 
+            description="This will be saved",
+            embedding_model=embedding_model,
+            dim=dim
+        )
+        assert result1["status"] == "success"
+        
+        # 创建第二个collection
+        result2 = memory_service.create_collection(
+            name="unsaved_collection",
+            backend_type="VDB",
+            description="This will NOT be saved",
+            embedding_model=embedding_model,
+            dim=dim
+        )
+        assert result2["status"] == "success"
+        
+        # 向两个collection插入数据
+        for i in range(3):
+            # 保存的collection
+            memory_service.insert_data(
+                collection_name="saved_collection",
+                text=f"Saved document {i}",
+                metadata={"type": "saved", "index": i}
+            )
+            # 未保存的collection  
+            memory_service.insert_data(
+                collection_name="unsaved_collection", 
+                text=f"Unsaved document {i}",
+                metadata={"type": "unsaved", "index": i}
+            )
+        
+        # 只保存第一个collection
+        print("\n💾 Phase 2: Storing only the first collection...")
+        store_result = memory_service.store_collection("saved_collection")
+        assert store_result["status"] == "success"
+        print(f"✅ Stored 'saved_collection': {store_result}")
+        
+        # 验证两个collection都还在内存中
+        list_result = memory_service.list_collections()
+        assert list_result["status"] == "success"
+        assert len(list_result["collections"]) == 2
+        print(f"✅ Both collections still in memory: {[c['name'] for c in list_result['collections']]}")
+        
+        # 释放当前service实例
+        print("\n🔄 Phase 3: Creating new service instance to simulate restart...")
+        del memory_service
+        
+        # 重新创建服务实例 (模拟重启后从磁盘加载)
+        memory_service2 = MemoryService(data_dir=test_data_dir)
+        
+        # 验证数据丢失情况
+        print("\n🔍 Phase 4: Checking data after restart...")
+        
+        # 检查collections列表
+        list_result_after = memory_service2.list_collections()
+        assert list_result_after["status"] == "success"
+        collection_names = [c['name'] for c in list_result_after['collections']]
+        print(f"Collections after restart: {collection_names}")
+        
+        # 验证保存的collection仍然存在
+        if "saved_collection" in collection_names:
+            print("✅ 'saved_collection' found - data persisted correctly")
+            # 尝试检索数据验证完整性
+            retrieve_saved = memory_service2.retrieve_data(
+                collection_name="saved_collection",
+                query_text="Saved document",
+                topk=5,
+                with_metadata=True
+            )
+            if retrieve_saved["status"] == "success" and len(retrieve_saved["results"]) > 0:
+                print(f"✅ Saved collection data intact: {len(retrieve_saved['results'])} documents found")
+            else:
+                print("⚠️ Saved collection exists but data may be incomplete")
+        else:
+            print("❌ ERROR: 'saved_collection' not found after restart!")
+            
+        # 验证未保存的collection丢失
+        if "unsaved_collection" not in collection_names:
+            print("✅ 'unsaved_collection' correctly lost - demonstrates need for storage")
+        else:
+            print("⚠️ WARNING: 'unsaved_collection' unexpectedly persisted")
+            
+        print("\n🎯 Data loss test demonstrates:")
+        print("  - Collections must be explicitly stored to persist")
+        print("  - Only stored collections survive service restart") 
+        print("  - store_collection() provides selective persistence")
+        
+        return True
+        
+    except Exception as e:
+        print(f"❌ Data loss test failed: {e}")
+        traceback.print_exc()
+        return False
+        
+    finally:
+        # 清理测试数据
+        try:
+            test_data_dir = os.path.join(os.path.dirname(__file__), "test_data_loss")
+            if os.path.exists(test_data_dir):
+                shutil.rmtree(test_data_dir)
+            print("🧹 Data loss test cleanup completed")
+        except Exception as e:
+            print(f"⚠️ Data loss test cleanup error: {e}")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("SAGE Memory Service Integration Test")
+    print("=" * 60)
+    
+    # 运行基本功能测试
+    success = test_memory_service()
+    
+    if success:
+        print("\n🎉 Basic tests passed! Running data loss test...")
+        # 运行数据丢失测试
+        loss_test_success = test_data_loss_without_storage()
+        
+        if loss_test_success:
+            print("\n🎉 All tests passed! Memory service system is working correctly.")
+            print("✨ Storage functionality verified - data persistence works as expected.")
+        else:
+            print("\n💥 Data loss tests failed! Please check the logs above.")
+            exit(1)
+    else:
+        print("\n💥 Basic tests failed! Please check the logs above.")
+        exit(1)

sage/middleware/components/neuromem/utils/__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/utils/path_utils.py

@@ -0,0 +1,25 @@
+# import os
+
+# def get_default_data_dir():
+#     this_file = os.path.abspath(__file__)
+#     cur_dir = os.path.dirname(this_file)
+#     # 一直向上，直到找到 sage.middleware.services.neuromem. 目录
+#     while True:
+#         if os.path.basename(cur_dir) == "sage":
+#             project_root = os.path.dirname(cur_dir)
+#             data_dir = os.path.join(project_root, "data", "neuromem_data")
+#             os.makedirs(data_dir, exist_ok=True)
+#             return data_dir
+#         parent = os.path.dirname(cur_dir)
+#         if parent == cur_dir:
+#             raise FileNotFoundError("Could not find 'sage' directory in parent folders.")
+#         cur_dir = parent
+        
+import os
+
+def get_default_data_dir():
+    # 获取当前执行程序的工作目录
+    cur_dir = os.getcwd()
+    data_dir = os.path.join(cur_dir, "data", "neuromem_data")
+    os.makedirs(data_dir, exist_ok=True)
+    return data_dir

sage/middleware/components/sage_db/__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
+    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/sage_db/python/__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/examples/api_usage_tutorial.py

@@ -10,10 +10,10 @@ from typing import List, Dict, Any
 from sage.middleware.api import KVServiceAPI, VDBServiceAPI, MemoryServiceAPI, GraphServiceAPI
 
 # 导入具体服务实现和工厂函数
-from sage.middleware import (
+from sage.middleware.services import (
     KVService, VDBService, MemoryService, GraphService,
-    create_kv_service_factory, create_vdb_service_factory, 
-    create_memory_service_factory, create_graph_service_factory
+    create_kv_service_factory, create_vdb_service_factory,
+    create_memory_service_factory, create_graph_service_factory,
 )
 from sage.core.api.local_environment import LocalEnvironment
 

sage/middleware/examples/dag_microservices_demo.py

@@ -12,8 +12,8 @@ from sage.core.api.local_environment import LocalEnvironment
 from sage.middleware.services import (
     MemoryService,
     create_kv_service_factory,
-    create_vdb_service_factory, 
-    create_memory_service_factory
+    create_vdb_service_factory,
+    create_memory_service_factory,
 )
 
 
@@ -35,16 +35,15 @@ class SampleApplication:
             max_size=1000,
             ttl_seconds=3600  # 1小时过期
         )
-        self.env.register_service("kv_service", kv_factory.service_class, kv_factory)
+    self.env.register_service_factory("kv_service", kv_factory)
         
         # 注册VDB服务
         vdb_factory = create_vdb_service_factory(
             service_name="vdb_service",
-            collection_name="demo_vectors",
-            dimension=384,
-            persist_directory="./demo_vectors"
+            embedding_dimension=384,
+            index_type="IndexFlatL2",
         )
-        self.env.register_service("vdb_service", vdb_factory.service_class, vdb_factory)
+        self.env.register_service_factory("vdb_service", vdb_factory)
         
         # 注册Memory编排服务
         memory_factory = create_memory_service_factory(
@@ -52,7 +51,7 @@ class SampleApplication:
             kv_service_name="kv_service",
             vdb_service_name="vdb_service"
         )
-        self.env.register_service("memory_service", memory_factory.service_class, memory_factory)
+    self.env.register_service_factory("memory_service", memory_factory)
         
         print("✅ 所有服务已注册")
     

sage/middleware/examples/microservices_integration_demo.py

@@ -6,7 +6,7 @@ import time
 import numpy as np
 from typing import Dict, List, Any
 from sage.core.api.local_environment import LocalEnvironment
-from sage.middleware import (
+from sage.middleware.services import (
     create_kv_service_factory,
     create_vdb_service_factory,
     create_graph_service_factory, 
@@ -37,29 +37,27 @@ class SAGEMicroservicesDemo:
             max_size=100000,
             enable_text_index=True
         )
-        self.env.register_service("production_kv", kv_factory)
+    self.env.register_service_factory("production_kv", kv_factory)
         print("   ✅ KV Service ready - Key-Value storage + BM25 text search")
         
         # 2. VDB服务 - 向量数据库和语义搜索
         print("🔍 Setting up VDB Service...")
         vdb_factory = create_vdb_service_factory(
             service_name="production_vdb",
-            embedding_dimension=768,  # 常用的BERT embedding维度
+            embedding_dimension=768,
             index_type="IndexFlatL2",
-            metric="L2"
         )
-        self.env.register_service("production_vdb", vdb_factory)
+        self.env.register_service_factory("production_vdb", vdb_factory)
         print("   ✅ VDB Service ready - FAISS vector search")
         
         # 3. Graph服务 - 知识图谱和关系查询
         print("🕸️ Setting up Graph Service...")
         graph_factory = create_graph_service_factory(
             service_name="production_graph",
-            backend_type="memory",  # 生产环境可以用专门的图数据库
+            backend_type="memory",
             max_nodes=50000,
-            enable_inference=True
         )
-        self.env.register_service("production_graph", graph_factory)
+        self.env.register_service_factory("production_graph", graph_factory)
         print("   ✅ Graph Service ready - Knowledge graph + inference")
         
         # 4. Memory服务 - 智能记忆编排
@@ -67,12 +65,11 @@ class SAGEMicroservicesDemo:
         memory_factory = create_memory_service_factory(
             service_name="production_memory",
             kv_service_name="production_kv",
-            vdb_service_name="production_vdb", 
+            vdb_service_name="production_vdb",
             graph_service_name="production_graph",
             enable_knowledge_graph=True,
-            enable_cross_service_transactions=True
         )
-        self.env.register_service("production_memory", memory_factory)
+        self.env.register_service_factory("production_memory", memory_factory)
         print("   ✅ Memory Service ready - Intelligent memory orchestration")
         
         print("\n🎯 All microservices are now registered and ready!")

sage/middleware/examples/microservices_registration_demo.py

@@ -5,10 +5,10 @@ SAGE 微服务集成示例
 import time
 from sage.core.api.local_environment import LocalEnvironment
 from sage.middleware.services import (
-    create_kv_service_factory, 
+    create_kv_service_factory,
     create_vdb_service_factory,
     create_graph_service_factory,
-    create_memory_service_factory
+    create_memory_service_factory,
 )
 
 
@@ -52,20 +52,16 @@ def test_microservices_registration():
         max_size=50000,
         ttl_seconds=3600
     )
-    env2.register_service("redis_kv_service", kv_factory)
+    env2.register_service_factory("redis_kv_service", kv_factory)
     print("   ✅ Registered redis_kv_service")
     
     # VDB服务 - 使用ChromaDB后端
     vdb_factory = create_vdb_service_factory(
-        service_name="chroma_vdb_service", 
-        backend_type="memory",  # 可以改为"chroma"
-        # chroma_host="localhost",
-        # chroma_port=8000,
-        collection_name="sage_knowledge_base",
+        service_name="chroma_vdb_service",
         embedding_dimension=768,
-        distance_metric="cosine"
+        index_type="IndexFlatL2",
     )
-    env2.register_service("chroma_vdb_service", vdb_factory)
+    env2.register_service_factory("chroma_vdb_service", vdb_factory)
     print("   ✅ Registered chroma_vdb_service")
     
     # Graph服务 - 支持知识图谱
@@ -78,7 +74,7 @@ def test_microservices_registration():
         max_nodes=500000,
         max_relationships=2000000
     )
-    env2.register_service("knowledge_graph_service", graph_factory)
+    env2.register_service_factory("knowledge_graph_service", graph_factory)
     print("   ✅ Registered knowledge_graph_service")
     
     # Memory服务 - 协调所有服务
@@ -92,7 +88,7 @@ def test_microservices_registration():
         enable_caching=True,
         enable_knowledge_graph=True
     )
-    env2.register_service("unified_memory_service", memory_factory)
+    env2.register_service_factory("unified_memory_service", memory_factory)
     print("   ✅ Registered unified_memory_service")
     
     print("\n🎯 Service registration completed successfully!")