superlocalmemory 2.7.6 → 2.8.0

package/src/embeddings/__init__.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 SuperLocalMemory (superlocalmemory.com)
+"""embeddings package - Local Embedding Generation for SuperLocalMemory V2
+
+Re-exports all public classes and constants so that
+``from embeddings import EmbeddingEngine`` works.
+"""
+from embeddings.constants import (
+    MAX_BATCH_SIZE,
+    MAX_TEXT_LENGTH,
+    CACHE_MAX_SIZE,
+    SENTENCE_TRANSFORMERS_AVAILABLE,
+    SKLEARN_AVAILABLE,
+    TORCH_AVAILABLE,
+    CUDA_AVAILABLE,
+    MPS_AVAILABLE,
+    MEMORY_DIR,
+    EMBEDDING_CACHE_PATH,
+    MODEL_CACHE_PATH,
+)
+from embeddings.cache import LRUCache
+from embeddings.model_loader import ModelLoaderMixin
+from embeddings.engine import EmbeddingEngine
+from embeddings.database import add_embeddings_to_database
+from embeddings.cli import main
+
+__all__ = [
+    # Constants
+    "MAX_BATCH_SIZE",
+    "MAX_TEXT_LENGTH",
+    "CACHE_MAX_SIZE",
+    "SENTENCE_TRANSFORMERS_AVAILABLE",
+    "SKLEARN_AVAILABLE",
+    "TORCH_AVAILABLE",
+    "CUDA_AVAILABLE",
+    "MPS_AVAILABLE",
+    "MEMORY_DIR",
+    "EMBEDDING_CACHE_PATH",
+    "MODEL_CACHE_PATH",
+    # Classes
+    "LRUCache",
+    "ModelLoaderMixin",
+    "EmbeddingEngine",
+    # Functions
+    "add_embeddings_to_database",
+    "main",
+]

package/src/embeddings/cache.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 SuperLocalMemory (superlocalmemory.com)
+"""LRU cache for embedding vectors.
+"""
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+from collections import OrderedDict
+
+import numpy as np
+
+from embeddings.constants import CACHE_MAX_SIZE
+
+logger = logging.getLogger(__name__)
+
+
+class LRUCache:
+    """Simple LRU cache for embeddings."""
+
+    def __init__(self, max_size: int = CACHE_MAX_SIZE):
+        self.cache = OrderedDict()
+        self.max_size = max_size
+
+    def get(self, key: str) -> Optional[np.ndarray]:
+        """Get item from cache, moving to end (most recent)."""
+        if key not in self.cache:
+            return None
+
+        # Move to end (most recently used)
+        self.cache.move_to_end(key)
+        return np.array(self.cache[key])
+
+    def set(self, key: str, value: np.ndarray):
+        """Set item in cache, evicting oldest if full."""
+        if key in self.cache:
+            # Update existing
+            self.cache.move_to_end(key)
+            self.cache[key] = value.tolist()
+        else:
+            # Add new
+            if len(self.cache) >= self.max_size:
+                # Evict oldest
+                self.cache.popitem(last=False)
+            self.cache[key] = value.tolist()
+
+    def save(self, path: Path):
+        """Save cache to disk."""
+        try:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            with open(path, 'w') as f:
+                json.dump(dict(self.cache), f)
+            logger.debug(f"Saved {len(self.cache)} cached embeddings")
+        except Exception as e:
+            logger.error(f"Failed to save embedding cache: {e}")
+
+    def load(self, path: Path):
+        """Load cache from disk."""
+        if not path.exists():
+            return
+
+        try:
+            with open(path, 'r') as f:
+                data = json.load(f)
+                self.cache = OrderedDict(data)
+            logger.info(f"Loaded {len(self.cache)} cached embeddings")
+        except Exception as e:
+            logger.error(f"Failed to load embedding cache: {e}")
+            self.cache = OrderedDict()

package/src/embeddings/cli.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 SuperLocalMemory (superlocalmemory.com)
+"""CLI interface for testing the embedding engine.
+"""
+import sys
+import json
+import time
+import logging
+
+from embeddings.constants import MEMORY_DIR
+from embeddings.engine import EmbeddingEngine
+from embeddings.database import add_embeddings_to_database
+
+
+def main():
+    """Run the embedding engine CLI."""
+    # Configure logging
+    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
+
+    if len(sys.argv) < 2:
+        print("EmbeddingEngine CLI - Local Embedding Generation")
+        print("\nCommands:")
+        print("  python embedding_engine.py stats              # Show engine statistics")
+        print("  python embedding_engine.py generate           # Generate embeddings for database")
+        print("  python embedding_engine.py test               # Run performance test")
+        print("  python embedding_engine.py clear-cache        # Clear embedding cache")
+        sys.exit(0)
+
+    command = sys.argv[1]
+
+    if command == "stats":
+        engine = EmbeddingEngine()
+        stats = engine.get_stats()
+        print(json.dumps(stats, indent=2))
+
+    elif command == "generate":
+        db_path = MEMORY_DIR / "memory.db"
+        if not db_path.exists():
+            print(f"Database not found at {db_path}")
+            sys.exit(1)
+
+        print("Generating embeddings for all memories...")
+        engine = EmbeddingEngine()
+        add_embeddings_to_database(engine, db_path)
+        print("Generation complete!")
+        print(json.dumps(engine.get_stats(), indent=2))
+
+    elif command == "clear-cache":
+        engine = EmbeddingEngine()
+        engine.clear_cache()
+        engine.save_cache()
+        print("Cache cleared!")
+
+    elif command == "test":
+        print("Running embedding performance test...")
+
+        engine = EmbeddingEngine()
+
+        # Test single encoding
+        print("\nTest 1: Single text encoding")
+        text = "This is a test sentence for embedding generation."
+        start = time.time()
+        embedding = engine.encode(text)
+        elapsed = time.time() - start
+        print(f"  Time: {elapsed*1000:.2f}ms")
+        print(f"  Dimension: {len(embedding)}")
+        print(f"  Sample values: {embedding[:5]}")
+
+        # Test batch encoding
+        print("\nTest 2: Batch encoding (100 texts)")
+        texts = [f"This is test sentence number {i} with some content." for i in range(100)]
+        start = time.time()
+        embeddings = engine.encode(texts, batch_size=32)
+        elapsed = time.time() - start
+        print(f"  Time: {elapsed*1000:.2f}ms ({100/elapsed:.0f} texts/sec)")
+        print(f"  Shape: {embeddings.shape}")
+
+        # Test cache
+        print("\nTest 3: Cache performance")
+        start = time.time()
+        embedding_cached = engine.encode(text)
+        elapsed = time.time() - start
+        print(f"  Cache hit time: {elapsed*1000:.4f}ms")
+        print(f"  Speedup: {(elapsed*1000):.0f}x faster")
+
+        # Test similarity
+        print("\nTest 4: Similarity computation")
+        text1 = "The weather is nice today."
+        text2 = "It's a beautiful day outside."
+        text3 = "Python is a programming language."
+
+        emb1 = engine.encode(text1)
+        emb2 = engine.encode(text2)
+        emb3 = engine.encode(text3)
+
+        sim_12 = engine.similarity(emb1, emb2)
+        sim_13 = engine.similarity(emb1, emb3)
+
+        print(f"  Similarity (weather vs beautiful day): {sim_12:.3f}")
+        print(f"  Similarity (weather vs programming): {sim_13:.3f}")
+
+        # Print stats
+        print("\nEngine statistics:")
+        print(json.dumps(engine.get_stats(), indent=2))
+
+        # Save cache
+        engine.save_cache()
+        print("\nCache saved!")
+
+    else:
+        print(f"Unknown command: {command}")
+        print("Run without arguments to see available commands.")

package/src/embeddings/constants.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 SuperLocalMemory (superlocalmemory.com)
+"""Shared constants and feature-detection for the embeddings package.
+"""
+import logging
+from pathlib import Path
+
+# SECURITY: Embedding generation limits to prevent resource exhaustion
+MAX_BATCH_SIZE = 128
+MAX_TEXT_LENGTH = 10_000
+CACHE_MAX_SIZE = 10_000
+
+# Optional sentence-transformers dependency
+SENTENCE_TRANSFORMERS_AVAILABLE = False
+try:
+    from sentence_transformers import SentenceTransformer  # noqa: F401
+    SENTENCE_TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    SENTENCE_TRANSFORMERS_AVAILABLE = False
+
+# Fallback: TF-IDF vectorization
+SKLEARN_AVAILABLE = False
+try:
+    from sklearn.feature_extraction.text import TfidfVectorizer  # noqa: F401
+    SKLEARN_AVAILABLE = True
+except ImportError:
+    SKLEARN_AVAILABLE = False
+
+# GPU detection
+TORCH_AVAILABLE = False
+CUDA_AVAILABLE = False
+MPS_AVAILABLE = False  # Apple Silicon
+
+try:
+    import torch
+    TORCH_AVAILABLE = True
+    CUDA_AVAILABLE = torch.cuda.is_available()
+    MPS_AVAILABLE = hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
+except ImportError:
+    pass
+
+MEMORY_DIR = Path.home() / ".claude-memory"
+EMBEDDING_CACHE_PATH = MEMORY_DIR / "embedding_cache.json"
+MODEL_CACHE_PATH = MEMORY_DIR / "models"  # Local model storage
+
+logger = logging.getLogger(__name__)

package/src/embeddings/database.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 SuperLocalMemory (superlocalmemory.com)
+"""Database integration for batch embedding generation.
+"""
+import json
+import sqlite3
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def add_embeddings_to_database(
+    engine,
+    db_path: Path,
+    embedding_column: str = 'embedding',
+    batch_size: int = 32
+):
+    """
+    Generate embeddings for all memories in database.
+
+    Args:
+        engine: An EmbeddingEngine instance
+        db_path: Path to SQLite database
+        embedding_column: Column name to store embeddings
+        batch_size: Batch size for processing
+    """
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    try:
+        # Check if embedding column exists
+        cursor.execute("PRAGMA table_info(memories)")
+        columns = {row[1] for row in cursor.fetchall()}
+
+        if embedding_column not in columns:
+            # Add column
+            logger.info(f"Adding '{embedding_column}' column to database")
+            cursor.execute(f'ALTER TABLE memories ADD COLUMN {embedding_column} TEXT')
+            conn.commit()
+
+        # Get memories without embeddings
+        cursor.execute(f'''
+            SELECT id, content, summary
+            FROM memories
+            WHERE {embedding_column} IS NULL OR {embedding_column} = ''
+        ''')
+        rows = cursor.fetchall()
+
+        if not rows:
+            logger.info("All memories already have embeddings")
+            conn.close()
+            return
+
+        logger.info(f"Generating embeddings for {len(rows)} memories...")
+
+        # Process in batches
+        for i in range(0, len(rows), batch_size):
+            batch = rows[i:i+batch_size]
+            memory_ids = [row[0] for row in batch]
+
+            # Combine content and summary
+            texts = []
+            for row in batch:
+                content = row[1] or ""
+                summary = row[2] or ""
+                text = f"{content} {summary}".strip()
+                texts.append(text)
+
+            # Generate embeddings
+            embeddings = engine.encode(texts, batch_size=batch_size)
+
+            # Store in database
+            for mem_id, embedding in zip(memory_ids, embeddings):
+                embedding_json = json.dumps(embedding.tolist())
+                cursor.execute(
+                    f'UPDATE memories SET {embedding_column} = ? WHERE id = ?',
+                    (embedding_json, mem_id)
+                )
+
+            conn.commit()
+            logger.info(f"Processed {min(i+batch_size, len(rows))}/{len(rows)} memories")
+
+        # Save cache
+        engine.save_cache()
+
+        logger.info(f"Successfully generated embeddings for {len(rows)} memories")
+
+    finally:
+        conn.close()

package/src/embeddings/engine.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 SuperLocalMemory (superlocalmemory.com)
+"""EmbeddingEngine - Core encoding logic for local embedding generation.
+"""
+import hashlib
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from embeddings.constants import (
+    MAX_BATCH_SIZE,
+    MAX_TEXT_LENGTH,
+    CACHE_MAX_SIZE,
+    SENTENCE_TRANSFORMERS_AVAILABLE,
+    SKLEARN_AVAILABLE,
+    TORCH_AVAILABLE,
+    CUDA_AVAILABLE,
+    MPS_AVAILABLE,
+    EMBEDDING_CACHE_PATH,
+    MODEL_CACHE_PATH,
+)
+from embeddings.cache import LRUCache
+from embeddings.model_loader import ModelLoaderMixin
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingEngine(ModelLoaderMixin):
+    """
+    Local embedding generation using sentence-transformers.
+
+    Features:
+    - all-MiniLM-L6-v2 model (384 dimensions, 80MB, fast)
+    - Batch processing for efficiency (up to 128 texts)
+    - GPU acceleration (CUDA/MPS) with automatic detection
+    - LRU cache for repeated queries (10K entries)
+    - Graceful fallback to TF-IDF if dependencies unavailable
+
+    Performance:
+    - CPU: ~100 embeddings/sec
+    - GPU (CUDA): ~1000 embeddings/sec
+    - Apple Silicon (MPS): ~500 embeddings/sec
+    - Cache hit: ~0.001ms
+    """
+
+    def __init__(
+        self,
+        model_name: str = "all-MiniLM-L6-v2",
+        device: Optional[str] = None,
+        cache_path: Optional[Path] = None,
+        model_cache_path: Optional[Path] = None,
+        use_cache: bool = True
+    ):
+        """
+        Initialize embedding engine.
+
+        Args:
+            model_name: Sentence transformer model name
+            device: Device to use ('cuda', 'mps', 'cpu', or None for auto)
+            cache_path: Custom path for embedding cache
+            model_cache_path: Custom path for model storage
+            use_cache: Whether to use LRU cache
+        """
+        self.model_name = model_name
+        self.cache_path = cache_path or EMBEDDING_CACHE_PATH
+        self.model_cache_path = model_cache_path or MODEL_CACHE_PATH
+        self.use_cache = use_cache
+
+        # Auto-detect device
+        if device is None:
+            if CUDA_AVAILABLE:
+                device = 'cuda'
+                logger.info("Using CUDA GPU acceleration")
+            elif MPS_AVAILABLE:
+                device = 'mps'
+                logger.info("Using Apple Silicon (MPS) GPU acceleration")
+            else:
+                device = 'cpu'
+                logger.info("Using CPU (consider GPU for faster processing)")
+        self.device = device
+
+        # Initialize model
+        self.model = None
+        self.dimension = 384  # Default for all-MiniLM-L6-v2
+        self.use_transformers = SENTENCE_TRANSFORMERS_AVAILABLE
+
+        # Initialize cache
+        self.cache = LRUCache(max_size=CACHE_MAX_SIZE) if use_cache else None
+        if self.cache:
+            self.cache.load(self.cache_path)
+
+        # Fallback: TF-IDF vectorizer
+        self.tfidf_vectorizer = None
+        self.tfidf_fitted = False
+
+        # Load model (from ModelLoaderMixin)
+        self._load_model()
+
+    def _get_cache_key(self, text: str) -> str:
+        """Generate cache key for text."""
+        return hashlib.sha256(text.encode('utf-8')).hexdigest()[:32]
+
+    def encode(
+        self,
+        texts: Union[str, List[str]],
+        batch_size: int = 32,
+        show_progress: bool = False,
+        normalize: bool = True
+    ) -> np.ndarray:
+        """
+        Generate embeddings for text(s).
+
+        Args:
+            texts: Single text or list of texts
+            batch_size: Batch size for processing (max: 128)
+            show_progress: Show progress bar for large batches
+            normalize: Normalize embeddings to unit length
+
+        Returns:
+            Array of shape (n_texts, dimension) or (dimension,) for single text
+        """
+        single_input = isinstance(texts, str)
+        if single_input:
+            texts = [texts]
+
+        if len(texts) == 0:
+            return np.array([])
+
+        batch_size = min(batch_size, MAX_BATCH_SIZE)
+
+        # Validate text length
+        for i, text in enumerate(texts):
+            if not isinstance(text, str):
+                raise ValueError(f"Text at index {i} is not a string")
+            if len(text) > MAX_TEXT_LENGTH:
+                logger.warning(f"Text {i} truncated from {len(text)} to {MAX_TEXT_LENGTH} chars")
+                texts[i] = text[:MAX_TEXT_LENGTH]
+
+        # Check cache for hits
+        embeddings = []
+        uncached_texts = []
+        uncached_indices = []
+
+        if self.cache:
+            for i, text in enumerate(texts):
+                cache_key = self._get_cache_key(text)
+                cached = self.cache.get(cache_key)
+                if cached is not None:
+                    embeddings.append((i, cached))
+                else:
+                    uncached_texts.append(text)
+                    uncached_indices.append(i)
+        else:
+            uncached_texts = texts
+            uncached_indices = list(range(len(texts)))
+
+        # Generate embeddings for uncached texts
+        if uncached_texts:
+            if self.use_transformers and self.model:
+                uncached_embeddings = self._encode_transformer(
+                    uncached_texts, batch_size=batch_size, show_progress=show_progress
+                )
+            elif self.tfidf_vectorizer:
+                uncached_embeddings = self._encode_tfidf(uncached_texts)
+            else:
+                raise RuntimeError("No embedding method available")
+
+            for i, text, embedding in zip(uncached_indices, uncached_texts, uncached_embeddings):
+                if self.cache:
+                    cache_key = self._get_cache_key(text)
+                    self.cache.set(cache_key, embedding)
+                embeddings.append((i, embedding))
+
+        # Sort by original index and extract embeddings
+        embeddings.sort(key=lambda x: x[0])
+        result = np.array([emb for _, emb in embeddings])
+
+        if normalize and len(result) > 0:
+            norms = np.linalg.norm(result, axis=1, keepdims=True)
+            norms[norms == 0] = 1
+            result = result / norms
+
+        if single_input:
+            return result[0]
+        return result
+
+    def encode_batch(
+        self,
+        texts: List[str],
+        batch_size: int = 32,
+        show_progress: bool = True
+    ) -> np.ndarray:
+        """Convenience method for batch encoding with progress."""
+        return self.encode(texts, batch_size=batch_size, show_progress=show_progress)
+
+    def similarity(self, embedding1: np.ndarray, embedding2: np.ndarray) -> float:
+        """Compute cosine similarity between two embeddings. Returns [0, 1]."""
+        emb1 = embedding1 / (np.linalg.norm(embedding1) + 1e-8)
+        emb2 = embedding2 / (np.linalg.norm(embedding2) + 1e-8)
+        similarity = np.dot(emb1, emb2)
+        return float(max(0.0, min(1.0, similarity)))
+
+    def save_cache(self):
+        """Save embedding cache to disk."""
+        if self.cache:
+            self.cache.save(self.cache_path)
+
+    def clear_cache(self):
+        """Clear embedding cache."""
+        if self.cache:
+            self.cache.cache.clear()
+            logger.info("Cleared embedding cache")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get embedding engine statistics."""
+        return {
+            'sentence_transformers_available': SENTENCE_TRANSFORMERS_AVAILABLE,
+            'use_transformers': self.use_transformers,
+            'sklearn_available': SKLEARN_AVAILABLE,
+            'torch_available': TORCH_AVAILABLE,
+            'cuda_available': CUDA_AVAILABLE,
+            'mps_available': MPS_AVAILABLE,
+            'device': self.device,
+            'model_name': self.model_name,
+            'dimension': self.dimension,
+            'cache_enabled': self.cache is not None,
+            'cache_size': len(self.cache.cache) if self.cache else 0,
+            'cache_max_size': CACHE_MAX_SIZE,
+            'model_loaded': self.model is not None or self.tfidf_vectorizer is not None
+        }
+
+    def add_to_database(
+        self,
+        db_path,
+        embedding_column: str = 'embedding',
+        batch_size: int = 32
+    ):
+        """
+        Generate embeddings for all memories in database.
+
+        Delegates to :func:`embeddings.database.add_embeddings_to_database`.
+        """
+        from embeddings.database import add_embeddings_to_database
+        add_embeddings_to_database(self, db_path, embedding_column, batch_size)