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

sage/middleware/components/sage_db/sage_db.py

@@ -1,395 +0,0 @@
-"""
-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/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)