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

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

@@ -1,293 +0,0 @@
-"""
-测试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/sage_db/sage_db.py

@@ -1,187 +0,0 @@
-"""
-SAGE DB - High-performance vector database with FAISS backend
-
-This module provides a Python interface to the SAGE DB vector database,
-which supports efficient similarity search with metadata filtering.
-"""
-
-import numpy as np
-from typing import List, Dict, Any, Optional, Callable, Tuple, Union
-try:
-    from . import _sage_db
-except ImportError:
-    import _sage_db
-
-# Re-export C++ classes and enums
-IndexType = _sage_db.IndexType
-DistanceMetric = _sage_db.DistanceMetric
-QueryResult = _sage_db.QueryResult
-SearchParams = _sage_db.SearchParams
-DatabaseConfig = _sage_db.DatabaseConfig
-SageDBException = _sage_db.SageDBException
-
-class SageDB:
-    """
-    High-performance vector database with FAISS backend.
-    
-    Supports efficient similarity search, metadata filtering, and hybrid search.
-    """
-    
-    def __init__(self, dimension: int, 
-                 index_type: IndexType = IndexType.AUTO,
-                 metric: DistanceMetric = DistanceMetric.L2):
-        """
-        Initialize a new SageDB instance.
-        
-        Args:
-            dimension: Vector dimension
-            index_type: Index type for similarity search
-            metric: Distance metric
-        """
-        self._db = _sage_db.create_database(dimension, index_type, metric)
-    
-    @classmethod
-    def from_config(cls, config: DatabaseConfig):
-        """Create SageDB from configuration object."""
-        instance = cls.__new__(cls)
-        instance._db = _sage_db.create_database(config)
-        return instance
-    
-    def add(self, vector: Union[List[float], np.ndarray], 
-            metadata: Optional[Dict[str, str]] = None) -> int:
-        """Add a single vector with optional metadata."""
-        if isinstance(vector, np.ndarray):
-            vector = vector.tolist()
-        return self._db.add(vector, metadata or {})
-    
-    def add_batch(self, vectors: Union[List[List[float]], np.ndarray],
-                  metadata: Optional[List[Dict[str, str]]] = None) -> List[int]:
-        """Add multiple vectors with optional metadata."""
-        if isinstance(vectors, np.ndarray):
-            if len(vectors.shape) != 2:
-                raise ValueError("Vectors array must be 2-dimensional")
-            return _sage_db.add_numpy(self._db, vectors, metadata or [])
-        else:
-            return self._db.add_batch(vectors, metadata or [])
-    
-    def search(self, query: Union[List[float], np.ndarray], 
-               k: int = 10, include_metadata: bool = True) -> List[QueryResult]:
-        """Search for similar vectors."""
-        if isinstance(query, np.ndarray):
-            return _sage_db.search_numpy(self._db, query, SearchParams(k))
-        return self._db.search(query, k, include_metadata)
-    
-    def search_with_params(self, query: Union[List[float], np.ndarray],
-                          params: SearchParams) -> List[QueryResult]:
-        """Search with custom parameters."""
-        if isinstance(query, np.ndarray):
-            return _sage_db.search_numpy(self._db, query, params)
-        return self._db.search(query, params)
-    
-    def filtered_search(self, query: Union[List[float], np.ndarray],
-                       params: SearchParams,
-                       filter_fn: Callable[[Dict[str, str]], bool]) -> List[QueryResult]:
-        """Search with metadata filtering."""
-        if isinstance(query, np.ndarray):
-            query = query.tolist()
-        return self._db.filtered_search(query, params, filter_fn)
-    
-    def search_by_metadata(self, query: Union[List[float], np.ndarray],
-                          params: SearchParams,
-                          metadata_key: str,
-                          metadata_value: str) -> List[QueryResult]:
-        """Search with specific metadata constraint."""
-        if isinstance(query, np.ndarray):
-            query = query.tolist()
-        return self._db.query_engine().search_with_metadata(query, params, metadata_key, metadata_value)
-    
-    def hybrid_search(self, query: Union[List[float], np.ndarray],
-                     params: SearchParams,
-                     text_query: str = "",
-                     vector_weight: float = 0.7,
-                     text_weight: float = 0.3) -> List[QueryResult]:
-        """Hybrid vector and text search."""
-        if isinstance(query, np.ndarray):
-            query = query.tolist()
-        return self._db.query_engine().hybrid_search(query, params, text_query, vector_weight, text_weight)
-    
-    def build_index(self):
-        """Build/train the search index."""
-        self._db.build_index()
-    
-    def train_index(self, training_vectors: Optional[Union[List[List[float]], np.ndarray]] = None):
-        """Train the index with training data."""
-        if training_vectors is None:
-            self._db.train_index()
-        elif isinstance(training_vectors, np.ndarray):
-            training_list = [training_vectors[i].tolist() for i in range(training_vectors.shape[0])]
-            self._db.train_index(training_list)
-        else:
-            self._db.train_index(training_vectors)
-    
-    def is_trained(self) -> bool:
-        """Check if the index is trained."""
-        return self._db.is_trained()
-    
-    def set_metadata(self, vector_id: int, metadata: Dict[str, str]) -> bool:
-        """Set metadata for a vector."""
-        return self._db.set_metadata(vector_id, metadata)
-    
-    def get_metadata(self, vector_id: int) -> Optional[Dict[str, str]]:
-        """Get metadata for a vector."""
-        return self._db.get_metadata(vector_id)
-    
-    def find_by_metadata(self, key: str, value: str) -> List[int]:
-        """Find vectors by metadata key-value."""
-        return self._db.find_by_metadata(key, value)
-    
-    def save(self, filepath: str):
-        """Save database to disk."""
-        self._db.save(filepath)
-    
-    def load(self, filepath: str):
-        """Load database from disk."""
-        self._db.load(filepath)
-    
-    @property
-    def size(self) -> int:
-        """Number of vectors in the database."""
-        return self._db.size()
-    
-    @property
-    def dimension(self) -> int:
-        """Vector dimension."""
-        return self._db.dimension()
-    
-    @property
-    def index_type(self) -> IndexType:
-        """Index type."""
-        return self._db.index_type()
-    
-    def get_search_stats(self) -> Dict[str, Any]:
-        """Get last search statistics."""
-        stats = self._db.query_engine().get_last_search_stats()
-        return {
-            'total_candidates': stats.total_candidates,
-            'filtered_candidates': stats.filtered_candidates,
-            'final_results': stats.final_results,
-            'search_time_ms': stats.search_time_ms,
-            'filter_time_ms': stats.filter_time_ms,
-            'total_time_ms': stats.total_time_ms
-        }
-
-# Convenience functions
-def create_database(dimension: int, 
-                   index_type: IndexType = IndexType.AUTO,
-                   metric: DistanceMetric = DistanceMetric.L2) -> SageDB:
-    """Create a new SageDB instance."""
-    return SageDB(dimension, index_type, metric)
-
-def create_database_from_config(config: DatabaseConfig) -> SageDB:
-    """Create SageDB from configuration."""
-    return SageDB.from_config(config)
-
-__all__ = [
-    'SageDB', 'IndexType', 'DistanceMetric', 'QueryResult', 'SearchParams',
-    'DatabaseConfig', 'SageDBException', 'create_database', 'create_database_from_config'
-]

sage/middleware/components/sage_db/tests/test_python.py

@@ -1,144 +0,0 @@
-import sys
-import os
-import numpy as np
-
-# Add the build directory to the path
-sys.path.append('/home/shuhao/SAGE/sage_ext/sage_db/build')
-
-try:
-    import sage_db_py
-    from sage_db import SageDatabase
-    print("✅ Import successful!")
-except ImportError as e:
-    print(f"❌ Import failed: {e}")
-    print("Make sure to build the module first with ./build.sh")
-    sys.exit(1)
-
-def test_python_api():
-    """Test the Python API wrapper"""
-    print("\n🐍 Testing Python API...")
-    
-    # Create database
-    db = SageDatabase(dimension=64)
-    
-    # Add single vector
-    vector = np.random.random(64).astype(np.float32)
-    metadata = {"text": "test vector", "category": "example"}
-    
-    vector_id = db.add_vector(vector, metadata)
-    print(f"Added vector with ID: {vector_id}")
-    
-    # Add batch
-    batch_vectors = np.random.random((10, 64)).astype(np.float32)
-    batch_metadata = [{"text": f"vector_{i}", "batch": "test"} for i in range(10)]
-    
-    batch_ids = db.add_batch(batch_vectors, batch_metadata)
-    print(f"Added batch of {len(batch_ids)} vectors")
-    
-    # Search
-    query = np.random.random(64).astype(np.float32)
-    results = db.search(query, k=3)
-    
-    print(f"Search returned {len(results)} results:")
-    for i, result in enumerate(results):
-        print(f"  {i+1}. ID: {result['id']}, Distance: {result['distance']:.4f}")
-        print(f"      Metadata: {result['metadata']}")
-    
-    # Metadata search
-    category_results = db.find_by_metadata("category", "example")
-    print(f"Found {len(category_results)} vectors with category='example'")
-    
-    print("✅ Python API test passed!")
-
-def test_pybind11_interface():
-    """Test the direct pybind11 interface"""
-    print("\n🔧 Testing pybind11 interface...")
-    
-    # Test vector store
-    vector_store = sage_db_py.VectorStore(32)
-    
-    # Add vectors
-    vec1 = np.array([1.0] * 32, dtype=np.float32)
-    vec2 = np.array([2.0] * 32, dtype=np.float32)
-    
-    id1 = vector_store.add(vec1)
-    id2 = vector_store.add(vec2)
-    
-    print(f"Added vectors with IDs: {id1}, {id2}")
-    print(f"Vector store size: {vector_store.size()}")
-    
-    # Search
-    query = np.array([1.1] * 32, dtype=np.float32)
-    results = vector_store.search(query, 2)
-    
-    print(f"Search results: {len(results)} items")
-    for result in results:
-        print(f"  ID: {result['id']}, Distance: {result['distance']:.4f}")
-    
-    # Test metadata store
-    metadata_store = sage_db_py.MetadataStore()
-    
-    metadata_store.set(id1, {"name": "first", "type": "test"})
-    metadata_store.set(id2, {"name": "second", "type": "test"})
-    
-    meta1 = metadata_store.get(id1)
-    print(f"Metadata for ID {id1}: {meta1}")
-    
-    type_matches = metadata_store.find_by_key_value("type", "test")
-    print(f"Found {len(type_matches)} vectors with type='test'")
-    
-    print("✅ pybind11 interface test passed!")
-
-def benchmark_python():
-    """Simple performance test"""
-    print("\n📊 Python performance benchmark...")
-    
-    import time
-    
-    dimension = 128
-    num_vectors = 1000
-    
-    db = SageDatabase(dimension=dimension)
-    
-    # Generate random data
-    vectors = np.random.random((num_vectors, dimension)).astype(np.float32)
-    metadata = [{"id": str(i), "batch": "benchmark"} for i in range(num_vectors)]
-    
-    # Time batch addition
-    start_time = time.time()
-    ids = db.add_batch(vectors, metadata)
-    add_time = time.time() - start_time
-    
-    print(f"Added {num_vectors} vectors in {add_time:.3f} seconds")
-    print(f"Rate: {num_vectors / add_time:.1f} vectors/second")
-    
-    # Time searches
-    num_queries = 100
-    query_vectors = np.random.random((num_queries, dimension)).astype(np.float32)
-    
-    start_time = time.time()
-    for query in query_vectors:
-        results = db.search(query, k=10)
-    search_time = time.time() - start_time
-    
-    print(f"Performed {num_queries} searches in {search_time:.3f} seconds")
-    print(f"Rate: {num_queries / search_time:.1f} searches/second")
-    
-    print("✅ Performance benchmark completed!")
-
-if __name__ == "__main__":
-    print("🧪 SAGE DB Python Test Suite")
-    print("============================")
-    
-    try:
-        test_python_api()
-        test_pybind11_interface()
-        benchmark_python()
-        
-        print("\n🎉 All Python tests passed!")
-        
-    except Exception as e:
-        print(f"\n❌ Test failed: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)