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

sage/middleware/components/neuromem/tests/__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/tests/core_test/__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.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/tests/core_test/collection_test/__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.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/tests/core_test/collection_test/kv_collection_test.py

@@ -0,0 +1,60 @@
+import pytest
+from sage.middleware.components.neuromem.memory_collection.kv_collection import KVMemoryCollection
+from sage.middleware.components.neuromem.memory_collection.base_collection import get_default_data_dir
+
+@pytest.fixture
+def setup_kv_collection():
+    col = KVMemoryCollection(name="demo")
+
+    col.insert("Hello Jack.", {})
+    col.insert("Hello Tom.", {})
+    col.insert("Hello Alice.", {})
+    col.insert("Jack and Tom say hi.", {})
+    col.insert("Alice in Wonderland.", {})
+    col.insert("Jacky is not Jack.", {})
+
+    col.create_index("global_index")
+    return col
+
+
+def test_kv_retrieve(setup_kv_collection):
+    col = setup_kv_collection
+
+    res1 = col.retrieve("Jack", topk=3, index_name="global_index")
+    assert set(res1) == {"Hello Jack.", "Jack and Tom say hi.", "Jacky is not Jack."}
+
+    res2 = col.retrieve("Alice", topk=2, index_name="global_index")
+    assert set(res2) == {"Hello Alice.", "Alice in Wonderland."}
+
+    res3 = col.retrieve("Tom", topk=2, index_name="global_index")
+    assert set(res3) == {"Hello Tom.", "Jack and Tom say hi."}
+
+
+def test_kv_update_and_delete(setup_kv_collection):
+    col = setup_kv_collection
+
+    col.delete("Hello Tom.")
+    res4 = col.retrieve("Tom", index_name="global_index")
+    assert "Hello Tom." not in res4
+
+    col.update("Hello Jack.", "Hello Jacky.", {}, "global_index")
+    res5 = col.retrieve("Jacky", index_name="global_index")
+
+    assert {"Jacky is not Jack.", "Hello Jacky."}.issubset(set(res5))
+
+
+def test_kv_persistence(setup_kv_collection):
+    col = setup_kv_collection
+    store_path = get_default_data_dir()
+    col_name = "demo"
+    col.store(store_path)
+
+    del col
+
+    col2 = KVMemoryCollection.load(col_name)
+    res = col2.retrieve("Jacky", index_name="global_index")
+
+
+    assert "Jacky is not Jack." in res
+
+    KVMemoryCollection.clear(col_name)

sage/middleware/components/neuromem/tests/core_test/collection_test/vdb_collection_test.py

@@ -0,0 +1,88 @@
+# file sage/sage_tests/neuromem_test/core_test/collection_test/vdb_collection_test.py
+
+import os
+import time
+import pytest
+from sage.middleware.components.neuromem.memory_collection.base_collection import get_default_data_dir
+from sage.utils.embedding_methods.mockembedder import MockTextEmbedder
+from sage.middleware.components.neuromem.memory_collection.vdb_collection import VDBMemoryCollection
+
+
+def almost_equal_dict(d1, d2, float_tol=1e-3):
+    # 只对所有值都是float的dict做容忍，否则严格等价
+    if d1.keys() != d2.keys():
+        return False
+    for k in d1:
+        v1, v2 = d1[k], d2[k]
+        if isinstance(v1, float) and isinstance(v2, float):
+            if abs(v1 - v2) > float_tol:
+                return False
+        else:
+            if v1 != v2:
+                return False
+    return True
+
+@pytest.fixture
+def setup_vdb_collection():
+    default_model = MockTextEmbedder(fixed_dim=16)
+    col = VDBMemoryCollection("vdb_demo", default_model, 16)
+    col.add_metadata_field("source")
+    col.add_metadata_field("lang")
+    col.add_metadata_field("timestamp")
+
+    current_time = time.time()
+    texts = [
+        ("hello world", {"source": "user", "lang": "en", "timestamp": current_time - 3600}),
+        ("你好，世界", {"source": "user", "lang": "zh", "timestamp": current_time - 1800}),
+        ("bonjour le monde", {"source": "web", "lang": "fr", "timestamp": current_time}),
+    ]
+    for t, meta in texts:
+        col.insert(t, meta)
+
+    col.create_index("global_index")
+    col.create_index("en_index", metadata_filter_func=lambda m: m.get("lang") == "en", description="English only")
+    col.create_index("user_index", metadata_filter_func=lambda m: m.get("source") == "user")
+    return col, current_time
+
+def test_vdb_persist(setup_vdb_collection):
+    col, current_time = setup_vdb_collection
+
+    # 检索校验
+    res = col.retrieve("hello", topk=3, index_name="global_index")
+    assert set(res) & {"hello world"} == {"hello world"}
+
+    res = col.retrieve("hello", index_name="en_index")
+    assert set(res) & {"hello world"} == {"hello world"}
+
+    res = col.retrieve("你好", index_name="user_index")
+    assert set(res) & {"你好，世界"} == {"你好，世界"}
+
+    # 持久化保存
+    store_path = get_default_data_dir()
+    col_name = "vdb_demo"
+    col.store(store_path)
+
+    # 清除内存对象
+    del col
+
+    # 恢复对象并回归测试
+    default_model2 = MockTextEmbedder(fixed_dim=16)
+    col2 = VDBMemoryCollection.load(col_name, embedding_model=default_model2)
+
+    # 再检索
+    res = col2.retrieve("hello", index_name="global_index")
+    assert set(res) & {"hello world"} == {"hello world"}
+
+    res = col2.retrieve("你好", index_name="user_index")
+    assert set(res) & {"你好，世界"} == {"你好，世界"}
+
+    # 校验metadata一致性
+    meta = col2.metadata_storage.get(col2._get_stable_id("hello world"))
+    assert almost_equal_dict(meta, {"source": "user", "lang": "en", "timestamp": current_time - 3600})
+
+    # 校验索引条件
+    idx_meta = col2.indexes["en_index"]
+    assert idx_meta.get("description", "") == "English only"
+
+    # 删除磁盘数据
+    VDBMemoryCollection.clear(col_name)

sage/middleware/components/neuromem/tests/core_test/manager_test.py

@@ -0,0 +1,154 @@
+# file sage/sage_tests/neuromem_test/core_test/manager_test.py
+# python -m sage.sage_tests.neuromem_test.core_test.manager_test
+
+from time import sleep
+
+
+def test_vdb(do_reload=True, do_delete=True):
+    from sage.middleware.components.neuromem.memory_manager import MemoryManager
+    from sage.utils.embedding_methods.mockembedder import MockTextEmbedder
+    import os
+
+    def colored(text, color):
+        colors = {"green": "\033[92m", "red": "\033[91m", "yellow": "\033[93m", "reset": "\033[0m"}
+        return colors.get(color, "") + text + colors["reset"]
+        # return text  # 自动化测试不需要彩色输出
+
+    def print_test_case(desc, expected, actual):
+        status = "通过" if expected == actual or (isinstance(expected, set) and set(expected) == set(actual)) else "不通过"
+        print(f"【{desc}】")
+        print(f"预期结果：{expected}")
+        print(f"实际结果：{actual}")
+        print(f"测试情况：{status}\n")
+
+    # 1. 创建 MemoryManager 和 VDBCollection
+    mgr = MemoryManager()
+    embedder = MockTextEmbedder(fixed_dim=16)
+    col = mgr.create_collection(name="test_vdb", backend_type="VDB", description="测试VDB集合", embedding_model=embedder, dim=16)
+    col.add_metadata_field("tag")  # 注册元数据字段tag
+
+    # 2. 插入数据
+    col.insert("Alpha", {"tag": "A"})
+    col.insert("Beta", {"tag": "B"})
+    col.insert("Gamma", {"tag": "A"})
+
+    # 3. 创建索引
+    col.create_index("global_index")
+    col.create_index("tag_A_index", metadata_filter_func=lambda m: m.get("tag") == "A")
+
+    # 4. 检索测试
+    res1 = col.retrieve("Alpha", topk=1, index_name="global_index")
+    print_test_case("检索Alpha", ["Alpha"], res1)
+
+    res2 = set(col.retrieve("Alpha", topk=5, index_name="tag_A_index"))  # Alpha 的向量检索，tag为A
+    print_test_case("tag_A_index 检索", {'Gamma', 'Alpha', 'Beta'}, res2)
+
+    # 5. 持久化
+    mgr.store_collection()
+    print("数据已保存到磁盘！")
+    data_dir = mgr.data_dir
+    print("目录为：", os.path.join(data_dir, "vdb_collection", "test_vdb"))
+
+    # 6. 清空对象
+    del mgr, col
+    print("内存对象已清除。")
+
+    # 7. 读取持久化数据，自动测试时直接执行
+    if do_reload:
+    from sage.middleware.components.neuromem.memory_manager import MemoryManager
+        embedder2 = MockTextEmbedder(fixed_dim=16)
+        mgr2 = MemoryManager()
+        col2 = mgr2.connect_collection("test_vdb", embedding_model=embedder2)
+        res3 = set(col2.retrieve("Alpha", topk=5, index_name="tag_A_index"))
+        print_test_case("恢复后tag_A_index 检索", {'Gamma', 'Alpha', 'Beta'}, res3)
+    else:
+        print("跳过加载测试。")
+
+    sleep(10)  # 等待一会儿，确保所有操作完成
+    # 8. 删除所有数据
+    if do_delete:
+    from sage.middleware.components.neuromem.memory_collection.vdb_collection import VDBMemoryCollection
+        VDBMemoryCollection.clear("test_vdb", data_dir)
+        manager_json = os.path.join(data_dir, "manager.json")
+        if os.path.exists(manager_json):
+            os.remove(manager_json)
+        print("所有数据已删除！")
+    else:
+        print("未执行删除。")
+
+# 作为测试脚本被自动调用
+if __name__ == "__main__":
+    test_vdb()
+
+
+# def kvtest():
+#     from sage.core.sage.middleware.services.neuromem..memory_manager import MemoryManager
+#     import os
+
+#     def colored(text, color):
+#         colors = {"green": "\033[92m", "red": "\033[91m", "yellow": "\033[93m", "reset": "\033[0m"}
+#         return colors.get(color, "") + text + colors["reset"]
+
+#     def print_test_case(desc, expected, actual):
+#         status = "通过" if expected == actual or (isinstance(expected, set) and set(expected) == set(actual)) else "不通过"
+#         color = "green" if status == "通过" else "red"
+#         print(f"【{desc}】")
+#         print(f"预期结果：{expected}")
+#         print(f"实际结果：{actual}")
+#         print(f"测试情况：{colored(status, color)}\n")
+
+#     # 1. 创建 MemoryManager 和 KVCollection
+#     mgr = MemoryManager()
+#     col = mgr.create_collection(name="test_kv", backend_type="KV", description="测试KV集合")
+#     col.add_metadata_field("tag")  # 注册元数据字段tag
+
+#     # 2. 插入数据
+#     col.insert("Alpha", {"tag": "A"})
+#     col.insert("Beta", {"tag": "B"})
+#     col.insert("Gamma", {"tag": "A"})
+
+#     # 3. 创建索引
+#     col.create_index("global_index")
+#     col.create_index("tag_A_index", metadata_filter_func=lambda m: m.get("tag") == "A")
+
+#     # 4. 检索测试
+#     res1 = col.retrieve("Alpha", topk=1, index_name="global_index")
+#     print_test_case("检索Alpha", ["Alpha"], res1)
+
+#     res2 = set(col.retrieve("a", topk=5, index_name="tag_A_index"))  # 检索包含'a'的文本，tag为A
+#     print_test_case("tag_A_index 检索", {"Alpha", "Gamma"}, res2)
+
+#     # 5. 持久化
+#     mgr.store_collection()
+#     print(colored("数据已保存到磁盘！", "yellow"))
+#     data_dir = mgr.data_dir
+#     print("目录为：", os.path.join(data_dir, "kv_collection", "test_kv"))
+
+#     # 6. 清空对象
+#     del mgr, col
+#     print(colored("内存对象已清除。", "yellow"))
+
+#     # 7. 读取持久化数据
+#     user_input = input(colored("输入 yes 加载刚才保存的数据: ", "yellow"))
+#     if user_input.strip().lower() == "yes":
+#         mgr2 = MemoryManager()  # 会自动加载 manager.json 和 collection
+#         col2 = mgr2.get_collection("test_kv")
+#         # 检查是否正常恢复
+#         res3 = set(col2.retrieve("a", topk=5, index_name="tag_A_index"))
+#         print_test_case("恢复后tag_A_index 检索", {"Alpha", "Gamma"}, res3)
+#     else:
+#         print(colored("跳过加载测试。", "yellow"))
+
+#     # 8. 删除所有数据
+#     user_input = input(colored("输入 yes 删除磁盘所有数据: ", "yellow"))
+#     if user_input.strip().lower() == "yes":
+#         from sage.core.sage.middleware.services.neuromem.memory_collection.kv_collection import KVMemoryCollection
+#         KVMemoryCollection.clear("test_kv", data_dir)
+#         manager_json = os.path.join(data_dir, "manager.json")
+#         if os.path.exists(manager_json):
+#             os.remove(manager_json)
+#         print(colored("所有数据已删除！", "green"))
+#     else:
+#         print(colored("未执行删除。", "yellow"))
+
+# kvtest()