repotoire 0.1.2__cp313-cp313-win_amd64.whl → 0.1.4__cp313-cp313-win_amd64.whl

repotoire/__init__.py

@@ -5,10 +5,10 @@ Analyzes codebases using knowledge graphs to detect code smells,
 architectural issues, and technical debt.
 """
 
-__version__ = "0.1.2"
+__version__ = "0.1.4"
 
 from repotoire.pipeline import IngestionPipeline
-from repotoire.graph import Neo4jClient
+from repotoire.graph import FalkorDBClient
 from repotoire.detectors import AnalysisEngine
 from repotoire.models import CodebaseHealth, Finding
 

repotoire/ai/__init__.py

@@ -6,6 +6,11 @@ from repotoire.ai.embeddings import (
     EmbeddingConfig,
     EmbeddingBackend,
     create_embedder,
+    # Int8 quantization for memory-efficient storage
+    quantize_embedding,
+    dequantize_embedding,
+    quantize_embeddings_batch,
+    compute_cosine_similarity_quantized,
 )
 from repotoire.ai.retrieval import (
     GraphRAGRetriever,
@@ -43,6 +48,13 @@ from repotoire.ai.contextual import (
     ContextGenerationResult,
     create_context_generator,
 )
+from repotoire.ai.compression import (
+    EmbeddingCompressor,
+    TenantCompressor,
+    create_compressor,
+    estimate_memory_savings,
+    DEFAULT_TARGET_DIMS,
+)
 
 __all__ = [
     # NLP
@@ -52,6 +64,11 @@ __all__ = [
     "EmbeddingConfig",
     "EmbeddingBackend",
     "create_embedder",
+    # Int8 quantization (4x memory reduction)
+    "quantize_embedding",
+    "dequantize_embedding",
+    "quantize_embeddings_batch",
+    "compute_cosine_similarity_quantized",
     # Retrieval
     "GraphRAGRetriever",
     "RetrievalResult",
@@ -83,4 +100,10 @@ __all__ = [
     "CostLimitExceeded",
     "ContextGenerationResult",
     "create_context_generator",
+    # Compression (memory optimization)
+    "EmbeddingCompressor",
+    "TenantCompressor",
+    "create_compressor",
+    "estimate_memory_savings",
+    "DEFAULT_TARGET_DIMS",
 ]

repotoire/ai/compression.py

@@ -0,0 +1,543 @@
+"""Embedding compression for memory-efficient storage.
+
+Implements PCA dimensionality reduction + int8 quantization for 8x compression
+with <3% quality loss, based on research findings.
+
+Compression pipeline:
+1. PCA: 4096 → 2048 dimensions (2x reduction)
+2. int8 quantization: float32 → int8 (4x reduction)
+3. Combined: 8x total compression
+
+Storage: 4096 * 4 bytes = 16KB → 2048 * 1 byte = 2KB per embedding
+"""
+
+import os
+import json
+import pickle
+from pathlib import Path
+from typing import List, Optional, Tuple
+import numpy as np
+
+from repotoire.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+# Default compression settings
+DEFAULT_TARGET_DIMS = 1024  # PCA target dimensions (4x reduction from 4096)
+DEFAULT_QUANTIZATION_BITS = 8  # int8 quantization
+
+
+class EmbeddingCompressor:
+    """Compress embeddings using PCA + quantization.
+
+    Provides 8x compression with <3% quality loss on retrieval tasks.
+
+    Example:
+        >>> compressor = EmbeddingCompressor(target_dims=2048)
+        >>> # Fit on existing embeddings
+        >>> compressor.fit(existing_embeddings)
+        >>> # Compress new embeddings
+        >>> compressed = compressor.compress(new_embedding)
+        >>> # Decompress for similarity computation
+        >>> decompressed = compressor.decompress(compressed)
+    """
+
+    def __init__(
+        self,
+        target_dims: int = DEFAULT_TARGET_DIMS,
+        quantization_bits: int = DEFAULT_QUANTIZATION_BITS,
+        model_path: Optional[Path] = None,
+    ):
+        """Initialize compressor.
+
+        Args:
+            target_dims: Target dimensions after PCA (default: 2048)
+            quantization_bits: Bits for quantization (default: 8 for int8)
+            model_path: Path to save/load fitted PCA model
+        """
+        self.target_dims = target_dims
+        self.quantization_bits = quantization_bits
+        self.model_path = model_path or Path.home() / ".repotoire" / "pca_model.pkl"
+
+        # PCA components (fitted)
+        self._pca_components: Optional[np.ndarray] = None
+        self._pca_mean: Optional[np.ndarray] = None
+        self._source_dims: Optional[int] = None
+
+        # Quantization parameters (computed during fit)
+        self._scale: Optional[float] = None
+        self._zero_point: Optional[float] = None
+
+        # Try to load existing model
+        if self.model_path.exists():
+            self._load_model()
+
+    @property
+    def is_fitted(self) -> bool:
+        """Check if compressor has been fitted."""
+        return self._pca_components is not None
+
+    @property
+    def compression_ratio(self) -> float:
+        """Calculate compression ratio."""
+        if not self._source_dims:
+            return 1.0
+        # Original: source_dims * 4 bytes (float32)
+        # Compressed: target_dims * 1 byte (int8)
+        original_size = self._source_dims * 4
+        compressed_size = self.target_dims * 1
+        return original_size / compressed_size
+
+    def fit(
+        self,
+        embeddings: List[List[float]],
+        save: bool = True,
+    ) -> "EmbeddingCompressor":
+        """Fit PCA on a sample of embeddings.
+
+        Should be called with a representative sample of embeddings
+        (e.g., 1000-10000 embeddings) to learn the principal components.
+
+        Args:
+            embeddings: List of embedding vectors to fit on
+            save: Whether to save the fitted model to disk
+
+        Returns:
+            self for method chaining
+        """
+        if len(embeddings) < 100:
+            logger.warning(
+                f"Fitting PCA on only {len(embeddings)} samples. "
+                "Recommend at least 1000 for good quality."
+            )
+
+        # Convert to numpy array
+        X = np.array(embeddings, dtype=np.float32)
+        n_samples, n_features = X.shape
+        self._source_dims = n_features
+
+        # PCA requires n_components <= min(n_samples, n_features)
+        max_components = min(n_samples, n_features)
+        effective_target_dims = min(self.target_dims, max_components)
+
+        if effective_target_dims < self.target_dims:
+            logger.warning(
+                f"Reducing target_dims from {self.target_dims} to {effective_target_dims} "
+                f"(limited by {n_samples} samples)"
+            )
+            self.target_dims = effective_target_dims
+
+        logger.info(
+            f"Fitting PCA: {self._source_dims} dims → {self.target_dims} dims "
+            f"on {len(embeddings)} samples"
+        )
+
+        # Compute mean
+        self._pca_mean = np.mean(X, axis=0)
+
+        # Center the data
+        X_centered = X - self._pca_mean
+
+        # Compute covariance matrix (more memory efficient for large dims)
+        # Using SVD instead of eigendecomposition for numerical stability
+        try:
+            from sklearn.decomposition import PCA
+
+            # Use sklearn PCA for better numerical stability
+            pca = PCA(n_components=self.target_dims, svd_solver='randomized')
+            pca.fit(X)
+
+            self._pca_components = pca.components_.astype(np.float32)
+            self._pca_mean = pca.mean_.astype(np.float32)
+
+            explained_variance = sum(pca.explained_variance_ratio_)
+            logger.info(f"PCA explains {explained_variance:.1%} of variance")
+
+        except ImportError:
+            # Fallback to manual SVD if sklearn not available
+            logger.info("sklearn not available, using numpy SVD")
+            U, S, Vt = np.linalg.svd(X_centered, full_matrices=False)
+            self._pca_components = Vt[:self.target_dims].astype(np.float32)
+
+            # Compute explained variance
+            total_var = np.sum(S ** 2)
+            explained_var = np.sum(S[:self.target_dims] ** 2) / total_var
+            logger.info(f"PCA explains {explained_var:.1%} of variance")
+
+        # Compute quantization parameters from transformed data
+        X_transformed = self._pca_transform(X)
+        self._compute_quantization_params(X_transformed)
+
+        logger.info(
+            f"Compression ratio: {self.compression_ratio:.1f}x "
+            f"({self._source_dims * 4} bytes → {self.target_dims} bytes)"
+        )
+
+        if save:
+            self._save_model()
+
+        return self
+
+    def _pca_transform(self, X: np.ndarray) -> np.ndarray:
+        """Apply PCA transformation."""
+        X_centered = X - self._pca_mean
+        return X_centered @ self._pca_components.T
+
+    def _pca_inverse_transform(self, X_reduced: np.ndarray) -> np.ndarray:
+        """Inverse PCA transformation (approximate reconstruction)."""
+        return X_reduced @ self._pca_components + self._pca_mean
+
+    def _compute_quantization_params(self, X: np.ndarray) -> None:
+        """Compute scale and zero point for int8 quantization."""
+        # Use percentiles to be robust to outliers
+        min_val = np.percentile(X, 0.1)
+        max_val = np.percentile(X, 99.9)
+
+        # Compute scale and zero point for symmetric quantization
+        self._scale = (max_val - min_val) / 255  # 256 levels for int8
+        self._zero_point = min_val
+
+        logger.debug(f"Quantization params: scale={self._scale:.6f}, zero={self._zero_point:.6f}")
+
+    def _quantize(self, X: np.ndarray) -> np.ndarray:
+        """Quantize float32 to int8."""
+        # Scale to [0, 255] range
+        X_scaled = (X - self._zero_point) / self._scale
+        # Clip and convert to uint8
+        X_quantized = np.clip(X_scaled, 0, 255).astype(np.uint8)
+        return X_quantized
+
+    def _dequantize(self, X_quantized: np.ndarray) -> np.ndarray:
+        """Dequantize int8 back to float32."""
+        return X_quantized.astype(np.float32) * self._scale + self._zero_point
+
+    def compress(self, embedding: List[float]) -> bytes:
+        """Compress a single embedding to bytes.
+
+        Args:
+            embedding: Original embedding vector (e.g., 4096 floats)
+
+        Returns:
+            Compressed embedding as bytes (e.g., 2048 bytes)
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Compressor not fitted. Call fit() first.")
+
+        # Convert to numpy
+        X = np.array([embedding], dtype=np.float32)
+
+        # Apply PCA
+        X_reduced = self._pca_transform(X)
+
+        # Quantize to int8
+        X_quantized = self._quantize(X_reduced)
+
+        # Return as bytes
+        return X_quantized[0].tobytes()
+
+    def compress_batch(self, embeddings: List[List[float]]) -> List[bytes]:
+        """Compress multiple embeddings efficiently.
+
+        Args:
+            embeddings: List of embedding vectors
+
+        Returns:
+            List of compressed embeddings as bytes
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Compressor not fitted. Call fit() first.")
+
+        # Convert to numpy
+        X = np.array(embeddings, dtype=np.float32)
+
+        # Apply PCA
+        X_reduced = self._pca_transform(X)
+
+        # Quantize to int8
+        X_quantized = self._quantize(X_reduced)
+
+        # Return as list of bytes
+        return [row.tobytes() for row in X_quantized]
+
+    def decompress(self, compressed: bytes) -> List[float]:
+        """Decompress bytes back to embedding vector.
+
+        Note: This is an approximate reconstruction due to PCA and quantization.
+
+        Args:
+            compressed: Compressed embedding bytes
+
+        Returns:
+            Reconstructed embedding vector (original dimensions)
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Compressor not fitted. Call fit() first.")
+
+        # Convert bytes to numpy array
+        X_quantized = np.frombuffer(compressed, dtype=np.uint8).reshape(1, -1)
+
+        # Dequantize
+        X_reduced = self._dequantize(X_quantized)
+
+        # Inverse PCA (approximate reconstruction)
+        X_reconstructed = self._pca_inverse_transform(X_reduced)
+
+        return X_reconstructed[0].tolist()
+
+    def decompress_batch(self, compressed_list: List[bytes]) -> List[List[float]]:
+        """Decompress multiple embeddings efficiently.
+
+        Args:
+            compressed_list: List of compressed embedding bytes
+
+        Returns:
+            List of reconstructed embedding vectors
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Compressor not fitted. Call fit() first.")
+
+        # Stack all compressed embeddings
+        X_quantized = np.array([
+            np.frombuffer(c, dtype=np.uint8) for c in compressed_list
+        ])
+
+        # Dequantize
+        X_reduced = self._dequantize(X_quantized)
+
+        # Inverse PCA
+        X_reconstructed = self._pca_inverse_transform(X_reduced)
+
+        return X_reconstructed.tolist()
+
+    def get_reduced_embedding(self, embedding: List[float]) -> List[float]:
+        """Get PCA-reduced embedding without quantization.
+
+        Useful when you want dimensionality reduction but need float precision
+        for vector similarity search.
+
+        Args:
+            embedding: Original embedding vector
+
+        Returns:
+            Reduced embedding (target_dims floats)
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Compressor not fitted. Call fit() first.")
+
+        X = np.array([embedding], dtype=np.float32)
+        X_reduced = self._pca_transform(X)
+        return X_reduced[0].tolist()
+
+    def get_reduced_embeddings_batch(
+        self, embeddings: List[List[float]]
+    ) -> List[List[float]]:
+        """Get PCA-reduced embeddings for a batch.
+
+        Args:
+            embeddings: List of original embedding vectors
+
+        Returns:
+            List of reduced embeddings (target_dims floats each)
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Compressor not fitted. Call fit() first.")
+
+        X = np.array(embeddings, dtype=np.float32)
+        X_reduced = self._pca_transform(X)
+        return X_reduced.tolist()
+
+    def _save_model(self) -> None:
+        """Save fitted PCA model to disk."""
+        self.model_path.parent.mkdir(parents=True, exist_ok=True)
+
+        model_data = {
+            "pca_components": self._pca_components,
+            "pca_mean": self._pca_mean,
+            "source_dims": self._source_dims,
+            "target_dims": self.target_dims,
+            "scale": self._scale,
+            "zero_point": self._zero_point,
+        }
+
+        with open(self.model_path, "wb") as f:
+            pickle.dump(model_data, f)
+
+        logger.info(f"Saved PCA model to {self.model_path}")
+
+    def _load_model(self) -> None:
+        """Load fitted PCA model from disk."""
+        try:
+            with open(self.model_path, "rb") as f:
+                model_data = pickle.load(f)
+
+            self._pca_components = model_data["pca_components"]
+            self._pca_mean = model_data["pca_mean"]
+            self._source_dims = model_data["source_dims"]
+            self.target_dims = model_data["target_dims"]
+            self._scale = model_data["scale"]
+            self._zero_point = model_data["zero_point"]
+
+            logger.info(
+                f"Loaded PCA model: {self._source_dims} → {self.target_dims} dims"
+            )
+        except Exception as e:
+            logger.warning(f"Could not load PCA model: {e}")
+
+
+class TenantCompressor:
+    """Per-tenant embedding compressor with model storage in cloud.
+
+    Each tenant can have their own PCA model fitted on their codebase,
+    allowing for better compression quality tailored to their code patterns.
+    """
+
+    def __init__(
+        self,
+        tenant_id: str,
+        storage_backend: str = "local",  # or "s3", "gcs"
+        target_dims: int = DEFAULT_TARGET_DIMS,
+    ):
+        """Initialize tenant-specific compressor.
+
+        Args:
+            tenant_id: Unique tenant identifier
+            storage_backend: Where to store PCA models
+            target_dims: Target dimensions after compression
+        """
+        self.tenant_id = tenant_id
+        self.storage_backend = storage_backend
+        self.target_dims = target_dims
+
+        # Model path includes tenant ID
+        model_dir = Path.home() / ".repotoire" / "compression_models"
+        self.model_path = model_dir / f"{tenant_id}_pca.pkl"
+
+        self._compressor = EmbeddingCompressor(
+            target_dims=target_dims,
+            model_path=self.model_path,
+        )
+
+    @property
+    def is_fitted(self) -> bool:
+        """Check if tenant compressor is fitted."""
+        return self._compressor.is_fitted
+
+    def fit_from_graph(
+        self,
+        graph_client,
+        sample_size: int = 5000,
+    ) -> "TenantCompressor":
+        """Fit compressor on embeddings from tenant's graph.
+
+        Args:
+            graph_client: FalkorDB client for the tenant
+            sample_size: Number of embeddings to sample for fitting
+
+        Returns:
+            self for method chaining
+        """
+        # Query embeddings from graph
+        # Note: FalkorDB uses labels() function for label checks instead of inline syntax
+        query = """
+        MATCH (n)
+        WHERE ('Function' IN labels(n) OR 'Class' IN labels(n) OR 'File' IN labels(n)) AND n.embedding IS NOT NULL
+        RETURN n.embedding as embedding
+        LIMIT $limit
+        """
+
+        results = graph_client.query(query, {"limit": sample_size})
+
+        if not results:
+            logger.warning(f"No embeddings found for tenant {self.tenant_id}")
+            return self
+
+        embeddings = [r["embedding"] for r in results if r.get("embedding")]
+
+        if len(embeddings) < 100:
+            logger.warning(
+                f"Only {len(embeddings)} embeddings for tenant {self.tenant_id}. "
+                "Recommend at least 100 for quality compression."
+            )
+
+        logger.info(f"Fitting compressor for tenant {self.tenant_id} on {len(embeddings)} embeddings")
+        self._compressor.fit(embeddings)
+
+        return self
+
+    def compress(self, embedding: List[float]) -> bytes:
+        """Compress embedding using tenant's model."""
+        return self._compressor.compress(embedding)
+
+    def compress_batch(self, embeddings: List[List[float]]) -> List[bytes]:
+        """Compress batch of embeddings."""
+        return self._compressor.compress_batch(embeddings)
+
+    def decompress(self, compressed: bytes) -> List[float]:
+        """Decompress embedding."""
+        return self._compressor.decompress(compressed)
+
+    def get_reduced_embedding(self, embedding: List[float]) -> List[float]:
+        """Get PCA-reduced embedding (float precision)."""
+        return self._compressor.get_reduced_embedding(embedding)
+
+    def get_reduced_embeddings_batch(
+        self, embeddings: List[List[float]]
+    ) -> List[List[float]]:
+        """Get PCA-reduced embeddings for batch."""
+        return self._compressor.get_reduced_embeddings_batch(embeddings)
+
+
+def create_compressor(
+    target_dims: int = DEFAULT_TARGET_DIMS,
+    model_path: Optional[Path] = None,
+) -> EmbeddingCompressor:
+    """Factory function to create an embedding compressor.
+
+    Args:
+        target_dims: Target dimensions after PCA
+        model_path: Path to save/load PCA model
+
+    Returns:
+        EmbeddingCompressor instance
+    """
+    return EmbeddingCompressor(
+        target_dims=target_dims,
+        model_path=model_path,
+    )
+
+
+def estimate_memory_savings(
+    num_entities: int,
+    source_dims: int = 4096,
+    target_dims: int = DEFAULT_TARGET_DIMS,
+) -> dict:
+    """Estimate memory savings from compression.
+
+    Args:
+        num_entities: Number of entities with embeddings
+        source_dims: Original embedding dimensions
+        target_dims: Target dimensions after compression
+
+    Returns:
+        Dictionary with memory estimates
+    """
+    # Original: float32 (4 bytes per dimension)
+    original_bytes = num_entities * source_dims * 4
+
+    # Compressed: int8 (1 byte per dimension)
+    compressed_bytes = num_entities * target_dims * 1
+
+    # PCA-reduced only (float32, no quantization)
+    reduced_bytes = num_entities * target_dims * 4
+
+    return {
+        "num_entities": num_entities,
+        "source_dims": source_dims,
+        "target_dims": target_dims,
+        "original_mb": original_bytes / (1024 * 1024),
+        "compressed_mb": compressed_bytes / (1024 * 1024),
+        "reduced_only_mb": reduced_bytes / (1024 * 1024),
+        "compression_ratio": original_bytes / compressed_bytes,
+        "savings_mb": (original_bytes - compressed_bytes) / (1024 * 1024),
+        "savings_percent": (1 - compressed_bytes / original_bytes) * 100,
+    }