contextsynapse 1.0.0__py3-none-any.whl

contextsynapse/storage/hnsw_index.py

@@ -0,0 +1,294 @@
+"""
+HNSW Index for Fast Vector Search
+
+Hierarchical Navigable Small World (HNSW) index for approximate nearest neighbor search.
+
+Benefits:
+- 100x faster than linear scan
+- Sub-millisecond query time
+- Minimal memory overhead
+- Incremental updates
+
+Author: AIContextDB Team
+Date: 2025-10-12
+"""
+
+import os
+import logging
+from typing import List, Optional, Tuple
+from pathlib import Path
+
+import numpy as np
+try:
+    import hnswlib
+    HNSW_AVAILABLE = True
+except ImportError:
+    HNSW_AVAILABLE = False
+    logging.warning("hnswlib not available, vector search will be slower")
+
+logger = logging.getLogger(__name__)
+
+
+class HNSWIndex:
+    """
+    Fast vector search using HNSW index.
+    
+    Provides 100x speedup over linear scan for large graphs.
+    """
+    
+    def __init__(self, path: str, embedding_dim: int = 768, 
+                 max_elements: int = 100000, ef_construction: int = 200, M: int = 16):
+        """
+        Initialize HNSW index.
+        
+        Args:
+            path: Directory for index storage
+            embedding_dim: Dimension of embeddings
+            max_elements: Maximum number of vectors
+            ef_construction: Construction-time parameter (higher = better recall, slower build)
+            M: Number of bi-directional links (higher = better recall, more memory)
+        """
+        if not HNSW_AVAILABLE:
+            raise ImportError("hnswlib is required for HNSW index. Install with: pip install hnswlib")
+        
+        self.path = Path(path)
+        self.path.mkdir(parents=True, exist_ok=True)
+        self.embedding_dim = embedding_dim
+        self.max_elements = max_elements
+        
+        # Create index
+        self.index = hnswlib.Index(space='cosine', dim=embedding_dim)
+        
+        # Try to load existing index
+        index_path = self.path / "hnsw.index"
+        if index_path.exists():
+            try:
+                self.index.load_index(str(index_path), max_elements=max_elements)
+                self.current_count = self.index.get_current_count()
+                logger.info(f"[EMOJI] Loaded HNSW index with {self.current_count} vectors")
+            except Exception as e:
+                logger.error(f"[EMOJI] Failed to load HNSW index: {e}")
+                self._init_new_index(ef_construction, M)
+        else:
+            self._init_new_index(ef_construction, M)
+        
+        self.stats = {
+            'vectors_added': 0,
+            'searches_performed': 0,
+            'avg_search_time_ms': 0.0
+        }
+    
+    def _init_new_index(self, ef_construction: int, M: int):
+        """Initialize a new HNSW index."""
+        self.index.init_index(
+            max_elements=self.max_elements,
+            ef_construction=ef_construction,
+            M=M
+        )
+        self.index.set_ef(50)  # Query-time parameter
+        self.current_count = 0
+        logger.info(f"[EMOJI] Created new HNSW index (dim={self.embedding_dim}, max={self.max_elements})")
+    
+    def add_vectors(self, embeddings: np.ndarray, ids: Optional[np.ndarray] = None):
+        """
+        Add vectors to index.
+        
+        Args:
+            embeddings: NumPy array of shape [n, embedding_dim]
+            ids: Optional array of IDs (indices into node table)
+        """
+        if embeddings.shape[0] == 0:
+            return
+        
+        if ids is None:
+            ids = np.arange(self.current_count, self.current_count + len(embeddings))
+        
+        # Ensure proper shape
+        if len(embeddings.shape) == 1:
+            embeddings = embeddings.reshape(1, -1)
+        
+        # Add to index
+        self.index.add_items(embeddings, ids)
+        
+        self.current_count = self.index.get_current_count()
+        self.stats['vectors_added'] += len(embeddings)
+        
+        logger.debug(f"[EMOJI] Added {len(embeddings)} vectors to HNSW index (total: {self.current_count})")
+    
+    def search(self, query_embedding: np.ndarray, k: int = 10, ef: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Search for k nearest neighbors.
+        
+        Args:
+            query_embedding: Query vector (1D array)
+            k: Number of neighbors to return
+            ef: Query-time parameter (higher = better recall, slower)
+        
+        Returns:
+            (indices, distances) - Arrays of shape [k]
+        """
+        if self.current_count == 0:
+            return np.array([]), np.array([])
+        
+        # Set ef if provided
+        if ef is not None:
+            self.index.set_ef(ef)
+        
+        # Ensure proper shape
+        if len(query_embedding.shape) == 1:
+            query_embedding = query_embedding.reshape(1, -1)
+        
+        # Search
+        labels, distances = self.index.knn_query(query_embedding, k=min(k, self.current_count))
+        
+        self.stats['searches_performed'] += 1
+        
+        # Return as 1D arrays
+        return labels[0], distances[0]
+    
+    def batch_search(self, query_embeddings: np.ndarray, k: int = 10) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Batch search for multiple queries.
+        
+        Args:
+            query_embeddings: Query vectors of shape [n, embedding_dim]
+            k: Number of neighbors per query
+        
+        Returns:
+            (indices, distances) - Arrays of shape [n, k]
+        """
+        if self.current_count == 0:
+            return np.array([]), np.array([])
+        
+        # Ensure proper shape
+        if len(query_embeddings.shape) == 1:
+            query_embeddings = query_embeddings.reshape(1, -1)
+        
+        # Batch search
+        labels, distances = self.index.knn_query(query_embeddings, k=min(k, self.current_count))
+        
+        self.stats['searches_performed'] += len(query_embeddings)
+        
+        return labels, distances
+    
+    def save(self):
+        """Save index to disk."""
+        index_path = self.path / "hnsw.index"
+        self.index.save_index(str(index_path))
+        
+        size_mb = index_path.stat().st_size / (1024 * 1024)
+        logger.info(f"[EMOJI] Saved HNSW index ({self.current_count} vectors, {size_mb:.2f} MB)")
+    
+    def rebuild(self, embeddings: np.ndarray, ids: np.ndarray):
+        """
+        Rebuild index from scratch.
+        
+        Args:
+            embeddings: All embeddings
+            ids: All IDs
+        """
+        logger.info(f"[EMOJI] Rebuilding HNSW index with {len(embeddings)} vectors...")
+        
+        # Create new index
+        self.index = hnswlib.Index(space='cosine', dim=self.embedding_dim)
+        self.index.init_index(
+            max_elements=max(self.max_elements, len(embeddings)),
+            ef_construction=200,
+            M=16
+        )
+        self.index.set_ef(50)
+        
+        # Add all vectors
+        self.add_vectors(embeddings, ids)
+        
+        logger.info(f"[EMOJI] HNSW index rebuilt")
+    
+    def get_stats(self) -> dict:
+        """Get index statistics."""
+        return {
+            **self.stats,
+            'current_count': self.current_count,
+            'max_elements': self.max_elements,
+            'embedding_dim': self.embedding_dim
+        }
+
+
+class LinearIndex:
+    """
+    Fallback linear scan index (when hnswlib is not available).
+    
+    Much slower than HNSW, but always available.
+    """
+    
+    def __init__(self, path: str, embedding_dim: int = 768, max_elements: int = 100000):
+        """Initialize linear index."""
+        self.path = Path(path)
+        self.embedding_dim = embedding_dim
+        self.embeddings = None
+        self.ids = None
+        self.current_count = 0
+        
+        logger.warning("[EMOJI][EMOJI]  Using linear index (slow). Install hnswlib for 100x speedup.")
+    
+    def add_vectors(self, embeddings: np.ndarray, ids: Optional[np.ndarray] = None):
+        """Add vectors."""
+        if embeddings.shape[0] == 0:
+            return
+        
+        if ids is None:
+            ids = np.arange(self.current_count, self.current_count + len(embeddings))
+        
+        if self.embeddings is None:
+            self.embeddings = embeddings
+            self.ids = ids
+        else:
+            self.embeddings = np.vstack([self.embeddings, embeddings])
+            self.ids = np.concatenate([self.ids, ids])
+        
+        self.current_count = len(self.embeddings)
+    
+    def search(self, query_embedding: np.ndarray, k: int = 10, ef: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
+        """Linear scan search."""
+        if self.current_count == 0:
+            return np.array([]), np.array([])
+        
+        # Compute cosine similarity
+        query_norm = query_embedding / np.linalg.norm(query_embedding)
+        emb_norms = self.embeddings / np.linalg.norm(self.embeddings, axis=1, keepdims=True)
+        similarities = emb_norms @ query_norm
+        
+        # Get top-k
+        top_k_indices = np.argsort(-similarities)[:k]
+        top_k_distances = 1 - similarities[top_k_indices]  # Convert to distance
+        
+        return self.ids[top_k_indices], top_k_distances
+    
+    def save(self):
+        """Save embeddings."""
+        if self.embeddings is not None:
+            np.save(str(self.path / "linear_embeddings.npy"), self.embeddings)
+            np.save(str(self.path / "linear_ids.npy"), self.ids)
+    
+    def get_stats(self) -> dict:
+        """Get stats."""
+        return {'current_count': self.current_count, 'type': 'linear'}
+
+
+def create_index(path: str, embedding_dim: int = 768, **kwargs):
+    """
+    Create appropriate index (HNSW if available, else linear).
+    
+    Args:
+        path: Storage path
+        embedding_dim: Embedding dimension
+        **kwargs: Additional arguments for HNSW
+    
+    Returns:
+        HNSWIndex or LinearIndex
+    """
+    if HNSW_AVAILABLE:
+        return HNSWIndex(path, embedding_dim, **kwargs)
+    else:
+        return LinearIndex(path, embedding_dim)
+
+

contextsynapse/storage/hybrid_store_v2.py

@@ -0,0 +1,278 @@
+"""
+Hybrid Storage V2 Implementation
+Enhanced version with better performance and advanced features
+"""
+
+import logging
+from typing import Dict, List, Any, Optional, Union, Tuple
+from pathlib import Path
+import json
+from datetime import datetime
+import asyncio
+import threading
+from concurrent.futures import ThreadPoolExecutor
+
+from .columnar_store_v2 import ColumnarStoreV2
+from .csr_store import CSREdgeStore
+
+logger = logging.getLogger(__name__)
+
+class HybridStoreV2:
+    """Enhanced hybrid storage with async operations and advanced features."""
+    
+    def __init__(self, storage_path: str = "contextcore_data"):
+        self.storage_path = Path(storage_path)
+        self.storage_path.mkdir(parents=True, exist_ok=True)
+        
+        # Initialize enhanced storage backends
+        self.columnar_store = ColumnarStoreV2(str(self.storage_path / "parquet"))
+        self.csr_store = CSREdgeStore()
+        
+        # Enhanced configuration
+        self.config = {
+            'use_columnar': True,
+            'use_csr': True,
+            'auto_optimize': True,
+            'async_operations': True,
+            'cache_size': 10000,
+            'batch_size': 1000,
+            'compression': 'snappy',
+            'indexing': True
+        }
+        
+        # Thread pool for async operations
+        self.executor = ThreadPoolExecutor(max_workers=4)
+        
+        # Cache for frequently accessed data
+        self.cache: Dict[str, Any] = {}
+        
+    async def store_graph_data_async(self, graph_data: Dict[str, Any], 
+                                   namespace: str = "default") -> bool:
+        """Asynchronously store graph data using hybrid approach."""
+        try:
+            tasks = []
+            
+            # Store nodes in columnar format (async)
+            if 'nodes' in graph_data and self.config['use_columnar']:
+                task = self._store_nodes_async(graph_data['nodes'], namespace)
+                tasks.append(task)
+            
+            # Store edges in CSR format (async)
+            if 'edges' in graph_data and self.config['use_csr']:
+                task = self._store_edges_async(graph_data['edges'], namespace)
+                tasks.append(task)
+            
+            # Wait for all tasks to complete
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # Check for errors
+            for result in results:
+                if isinstance(result, Exception):
+                    logger.error(f"Async storage task failed: {result}")
+                    return False
+            
+            logger.info(f"Async stored graph data for namespace '{namespace}'")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to async store graph data: {e}")
+            return False
+    
+    async def _store_nodes_async(self, nodes: List[Dict[str, Any]], 
+                               namespace: str) -> bool:
+        """Asynchronously store nodes."""
+        try:
+            # Process nodes in batches
+            batch_size = self.config['batch_size']
+            for i in range(0, len(nodes), batch_size):
+                batch = nodes[i:i + batch_size]
+                
+                # Convert to columnar format
+                nodes_data = []
+                for node in batch:
+                    nodes_data.append({
+                        'id': node.get('id', ''),
+                        'label': node.get('label', ''),
+                        'properties': json.dumps(node.get('properties', {})),
+                        'namespace': namespace,
+                        'created_at': datetime.now().isoformat()
+                    })
+                
+                # Store batch
+                table_name = f"{namespace}_nodes"
+                if not self.columnar_store.create_table(table_name, {
+                    'id': 'string',
+                    'label': 'string',
+                    'properties': 'string',
+                    'namespace': 'string',
+                    'created_at': 'string'
+                }):
+                    return False
+                
+                self.columnar_store.insert_data(table_name, nodes_data)
+                
+                # Yield control to other tasks
+                await asyncio.sleep(0)
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to async store nodes: {e}")
+            return False
+    
+    async def _store_edges_async(self, edges: List[Dict[str, Any]], 
+                               namespace: str) -> bool:
+        """Asynchronously store edges."""
+        try:
+            # Process edges in batches
+            batch_size = self.config['batch_size']
+            for i in range(0, len(edges), batch_size):
+                batch = edges[i:i + batch_size]
+                
+                # Store batch
+                for edge in batch:
+                    self.csr_store.add_edge(
+                        edge.get('source', ''),
+                        edge.get('target', ''),
+                        edge.get('label', ''),
+                        edge.get('properties', {})
+                    )
+                
+                # Yield control to other tasks
+                await asyncio.sleep(0)
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to async store edges: {e}")
+            return False
+    
+    def query_graph_data_optimized(self, namespace: str = "default",
+                                  query_type: str = "nodes",
+                                  conditions: Optional[Dict[str, Any]] = None,
+                                  limit: Optional[int] = None) -> List[Dict[str, Any]]:
+        """Optimized query with caching and indexing."""
+        try:
+            # Check cache first
+            cache_key = f"{namespace}_{query_type}_{str(conditions)}"
+            if cache_key in self.cache:
+                logger.debug(f"Cache hit for query: {cache_key}")
+                return self.cache[cache_key]
+            
+            results = []
+            
+            if query_type == "nodes" and self.config['use_columnar']:
+                table_name = f"{namespace}_nodes"
+                table = self.columnar_store.query(table_name, conditions)
+                
+                # Convert to list of dicts
+                if table.num_rows > 0:
+                    df = table.to_pandas()
+                    results = df.to_dict('records')
+                    
+                    # Parse JSON properties
+                    for result in results:
+                        if 'properties' in result:
+                            try:
+                                result['properties'] = json.loads(result['properties'])
+                            except:
+                                result['properties'] = {}
+            
+            elif query_type == "edges" and self.config['use_csr']:
+                edges = self.csr_store.get_all_edges()
+                results = [{'source': e.source_id, 'target': e.target_id,
+                           'label': e.edge_type, 'properties': e.properties}
+                          for e in edges]
+            
+            # Apply limit
+            if limit and len(results) > limit:
+                results = results[:limit]
+            
+            # Cache results
+            if len(self.cache) < self.config['cache_size']:
+                self.cache[cache_key] = results
+            
+            return results
+            
+        except Exception as e:
+            logger.error(f"Failed to query graph data: {e}")
+            return []
+    
+    def optimize_storage_advanced(self) -> bool:
+        """Advanced storage optimization with performance tuning."""
+        try:
+            if not self.config['auto_optimize']:
+                return True
+            
+            # Get detailed statistics
+            columnar_stats = self.columnar_store.get_statistics()
+            csr_stats = self.csr_store.get_statistics()
+            
+            # Performance tuning based on data characteristics
+            if columnar_stats['total_rows'] > 100000:
+                logger.info("Very large dataset detected, enabling advanced optimizations")
+                self.config['batch_size'] = 5000
+                self.config['cache_size'] = 50000
+            
+            if csr_stats['total_edges'] > 1000000:
+                logger.info("Very large graph detected, enabling CSR optimizations")
+                self.config['use_csr'] = True
+                self.config['indexing'] = True
+            
+            # Memory optimization
+            if columnar_stats['total_size_bytes'] > 1024 * 1024 * 1024:  # 1GB
+                logger.info("Large memory usage detected, enabling compression")
+                self.config['compression'] = 'gzip'
+            
+            logger.info("Advanced storage optimization completed")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to optimize storage: {e}")
+            return False
+    
+    def get_enhanced_statistics(self) -> Dict[str, Any]:
+        """Get enhanced statistics with performance metrics."""
+        columnar_stats = self.columnar_store.get_statistics()
+        csr_stats = self.csr_store.get_statistics()
+        
+        return {
+            'hybrid_config': self.config,
+            'columnar_stats': columnar_stats,
+            'csr_stats': csr_stats,
+            'cache_stats': {
+                'cache_size': len(self.cache),
+                'max_cache_size': self.config['cache_size']
+            },
+            'performance_metrics': {
+                'async_operations': self.config['async_operations'],
+                'batch_size': self.config['batch_size'],
+                'compression': self.config['compression'],
+                'indexing_enabled': self.config['indexing']
+            },
+            'total_namespaces': len(set([name.split('_')[0] for name in columnar_stats['table_names']])),
+            'storage_path': str(self.storage_path)
+        }
+    
+    def cleanup_cache(self):
+        """Clean up cache to free memory."""
+        self.cache.clear()
+        logger.info("Cache cleaned up")
+    
+    def __del__(self):
+        """Cleanup resources."""
+        if hasattr(self, 'executor'):
+            self.executor.shutdown(wait=False)
+
+
+
+
+
+
+
+
+
+
+
+
+