isage-middleware 0.1.0__py3-none-any.whl

sage/middleware/enterprise/sage_db/sage_db.py

@@ -0,0 +1,395 @@
+"""
+SAGE DB - 高性能向量数据库
+
+直接导出 C++ 扩展中的类和函数
+"""
+
+# 直接从 C++ 扩展导入所有必要的类
+try:
+    from ._sage_db import (
+        SageDB,
+        DatabaseConfig,
+        IndexType,
+        DistanceMetric,
+        QueryResult,
+        SearchParams,
+        SageDBException,
+        create_database,
+        index_type_to_string,
+        string_to_index_type,
+        distance_metric_to_string,
+        string_to_distance_metric
+    )
+    print("✓ SAGE DB C++ extension loaded successfully")
+except ImportError as e:
+    raise ImportError(f"Failed to import SAGE DB C++ extension: {e}")
+
+# 导出所有API
+__all__ = [
+    'SageDB',
+    'DatabaseConfig', 
+    'IndexType',
+    'DistanceMetric',
+    'QueryResult',
+    'SearchParams',
+    'SageDBException',
+    'create_database',
+    'index_type_to_string',
+    'string_to_index_type',
+    'distance_metric_to_string',
+    'string_to_distance_metric'
+]
+
+import numpy as np
+from typing import List, Dict, Any, Optional, Callable, Tuple, Union
+
+# 尝试导入C++扩展
+try:
+    from . import _sage_db
+    _cpp_available = True
+except ImportError:
+    try:
+        import _sage_db
+        _cpp_available = True
+    except ImportError:
+        _cpp_available = False
+        raise ImportError("SAGE DB C++ extension not available")
+
+# 重新导出C++类和枚举
+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 remove(self, vector_id: int) -> bool:
+        """Remove a vector by ID."""
+        return self._db.remove(vector_id)
+    
+    def update(self, vector_id: int, vector: Union[List[float], np.ndarray],
+               metadata: Optional[Dict[str, str]] = None) -> bool:
+        """Update a vector and its metadata."""
+        if isinstance(vector, np.ndarray):
+            vector = vector.tolist()
+        return self._db.update(vector_id, vector, metadata or {})
+    
+    def build_index(self):
+        """Build the search index."""
+        self._db.build_index()
+    
+    def train_index(self, training_data: Optional[List[List[float]]] = None):
+        """Train the index with training data."""
+        self._db.train_index(training_data or [])
+    
+    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."""
+        result = self._db.get_metadata(vector_id)
+        return result if result is not None else None
+    
+    def find_by_metadata(self, key: str, value: str) -> List[int]:
+        """Find vectors by metadata key-value pair."""
+        return self._db.find_by_metadata(key, value)
+    
+    def save(self, filepath: str):
+        """Save database to file."""
+        self._db.save(filepath)
+    
+    def load(self, filepath: str):
+        """Load database from file."""
+        self._db.load(filepath)
+    
+    def size(self) -> int:
+        """Get number of vectors in database."""
+        return self._db.size()
+    
+    def dimension(self) -> int:
+        """Get vector dimension."""
+        return self._db.dimension()
+    
+    def index_type(self) -> IndexType:
+        """Get index type."""
+        return self._db.index_type()
+    
+    def config(self) -> DatabaseConfig:
+        """Get database configuration."""
+        return self._db.config()
+    
+    # 高级功能访问
+    def query_engine(self):
+        """Get query engine for advanced search operations."""
+        return self._db.query_engine()
+    
+    def vector_store(self):
+        """Get vector store for low-level operations."""
+        return self._db.vector_store()
+    
+    def metadata_store(self):
+        """Get metadata store for metadata operations."""
+        return self._db.metadata_store()
+
+# 导出的API
+__all__ = [
+    'SageDB',
+    'IndexType',
+    'DistanceMetric', 
+    'QueryResult',
+    'SearchParams',
+    'DatabaseConfig',
+    'SageDBException'
+]
+
+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/enterprise/sage_db/tests/test_python.py

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