contextual-engine 0.1.0__py3-none-any.whl

contextual/embedding/cache.py

@@ -0,0 +1,455 @@
+"""Embedding cache for content-hash based deduplication.
+
+Caches embeddings keyed by content hash to avoid re-embedding unchanged chunks.
+Uses SQLite for persistence across indexing runs.
+
+Cache strategy:
+- Key: SHA-256 hash of chunk content (header + body)
+- Value: 768-dim embedding vector (JSON-encoded)
+- Lookup: O(1) with index on content_hash
+- Eviction: None (cache grows unbounded - embeddings are small)
+
+Performance impact:
+- Cache hit: ~0.1ms lookup vs ~55ms embedding generation (550x faster)
+- Storage: ~3KB per cached embedding (768 floats as JSON)
+- On 10K chunks: ~30MB cache, saves ~9 minutes of embedding time
+
+The cache table is separate from main schema to allow independent cleanup.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+from typing import TYPE_CHECKING
+
+from contextual.core.errors import ErrorCode, StorageError, storage_context
+from contextual.observability.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlite3 import Connection
+
+
+logger = get_logger(__name__)
+
+# ============================================================================
+# CACHE SCHEMA
+# ============================================================================
+
+CACHE_TABLE_DDL = """
+CREATE TABLE IF NOT EXISTS embedding_cache (
+    content_hash TEXT PRIMARY KEY,
+    embedding TEXT NOT NULL,
+    model_name TEXT NOT NULL,
+    created_at INTEGER NOT NULL,
+    hit_count INTEGER DEFAULT 0
+) STRICT;
+"""
+
+CACHE_INDEX_DDL = """
+CREATE INDEX IF NOT EXISTS idx_embedding_cache_model
+    ON embedding_cache(model_name, created_at);
+"""
+
+
+# ============================================================================
+# CACHE OPERATIONS
+# ============================================================================
+
+
+def initialize_cache(conn: Connection) -> None:
+    """Initialize embedding cache table and indexes.
+
+    Safe to call multiple times (uses IF NOT EXISTS).
+
+    Args:
+        conn: SQLite connection.
+
+    Raises:
+        StorageError: If cache initialization fails.
+    """
+    try:
+        cursor = conn.cursor()
+        cursor.execute(CACHE_TABLE_DDL)
+        cursor.execute(CACHE_INDEX_DDL)
+        conn.commit()
+
+        logger.debug("Embedding cache initialized")
+
+    except sqlite3.Error as e:
+        msg = f"Failed to initialize embedding cache: {e}"
+        raise StorageError(
+            code=ErrorCode.STORAGE_SQLITE_SCHEMA_INVALID,
+            message=msg,
+            context=storage_context(table="embedding_cache"),
+        ) from e
+
+
+def get_cached_embedding(
+    conn: Connection,
+    content_hash: str,
+    model_name: str,
+) -> list[float] | None:
+    """Retrieve cached embedding by content hash.
+
+    Increments hit_count on cache hit for analytics.
+
+    Args:
+        conn: SQLite connection.
+        content_hash: SHA-256 hash of content.
+        model_name: Embedding model identifier.
+
+    Returns:
+        768-dim embedding vector, or None if not cached.
+
+    Raises:
+        StorageError: If cache lookup fails.
+    """
+    try:
+        cursor = conn.cursor()
+
+        # Fetch cached embedding
+        cursor.execute(
+            """
+            SELECT embedding
+            FROM embedding_cache
+            WHERE content_hash = ? AND model_name = ?
+            """,
+            (content_hash, model_name),
+        )
+
+        row = cursor.fetchone()
+
+        if row is None:
+            return None
+
+        # Parse JSON embedding
+        embedding = json.loads(row[0])
+
+        # Increment hit count
+        cursor.execute(
+            """
+            UPDATE embedding_cache
+            SET hit_count = hit_count + 1
+            WHERE content_hash = ?
+            """,
+            (content_hash,),
+        )
+        conn.commit()
+
+        logger.debug(
+            "Cache hit",
+            content_hash=content_hash[:8],
+            model=model_name,
+        )
+
+        return embedding
+
+    except (sqlite3.Error, json.JSONDecodeError) as e:
+        msg = f"Failed to retrieve cached embedding: {e}"
+        raise StorageError(
+            code=ErrorCode.STORAGE_SQLITE_QUERY_FAILED,
+            message=msg,
+            context=storage_context(
+                table="embedding_cache",
+                query="get_cached_embedding",
+            ),
+        ) from e
+
+
+def put_cached_embedding(
+    conn: Connection,
+    content_hash: str,
+    embedding: list[float],
+    model_name: str,
+) -> None:
+    """Store embedding in cache.
+
+    Uses INSERT OR REPLACE for idempotent caching.
+
+    Args:
+        conn: SQLite connection.
+        content_hash: SHA-256 hash of content.
+        embedding: 768-dim embedding vector.
+        model_name: Embedding model identifier.
+
+    Raises:
+        StorageError: If cache write fails.
+    """
+    try:
+        import time
+
+        cursor = conn.cursor()
+        now_ms = int(time.time() * 1000)
+
+        # Serialize embedding as JSON
+        embedding_json = json.dumps(embedding)
+
+        # Insert or replace
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO embedding_cache
+                (content_hash, embedding, model_name, created_at, hit_count)
+            VALUES (?, ?, ?, ?, 0)
+            """,
+            (content_hash, embedding_json, model_name, now_ms),
+        )
+
+        conn.commit()
+
+        logger.debug(
+            "Cache write",
+            content_hash=content_hash[:8],
+            model=model_name,
+        )
+
+    except (sqlite3.Error, TypeError) as e:
+        msg = f"Failed to cache embedding: {e}"
+        raise StorageError(
+            code=ErrorCode.STORAGE_SQLITE_QUERY_FAILED,
+            message=msg,
+            context=storage_context(
+                table="embedding_cache",
+                query="put_cached_embedding",
+            ),
+        ) from e
+
+
+def get_cached_embeddings_batch(
+    conn: Connection,
+    content_hashes: list[str],
+    model_name: str,
+) -> dict[str, list[float]]:
+    """Retrieve multiple cached embeddings in one query.
+
+    More efficient than individual lookups for batches.
+
+    Args:
+        conn: SQLite connection.
+        content_hashes: List of content hashes to lookup.
+        model_name: Embedding model identifier.
+
+    Returns:
+        Dict mapping content_hash -> embedding for cache hits only.
+
+    Raises:
+        StorageError: If batch lookup fails.
+    """
+    if not content_hashes:
+        return {}
+
+    try:
+        cursor = conn.cursor()
+
+        # Build parameterized query
+        placeholders = ",".join("?" * len(content_hashes))
+        query = f"""
+            SELECT content_hash, embedding
+            FROM embedding_cache
+            WHERE content_hash IN ({placeholders})
+            AND model_name = ?
+        """
+
+        params = [*content_hashes, model_name]
+        cursor.execute(query, params)
+
+        # Parse results
+        results = {}
+        for content_hash, embedding_json in cursor.fetchall():
+            embedding = json.loads(embedding_json)
+            results[content_hash] = embedding
+
+        # Increment hit counts for all hits
+        if results:
+            hit_hashes = list(results.keys())
+            hit_placeholders = ",".join("?" * len(hit_hashes))
+            cursor.execute(
+                f"""
+                UPDATE embedding_cache
+                SET hit_count = hit_count + 1
+                WHERE content_hash IN ({hit_placeholders})
+                """,
+                hit_hashes,
+            )
+            conn.commit()
+
+        logger.debug(
+            "Batch cache lookup",
+            requested=len(content_hashes),
+            hits=len(results),
+            hit_rate=len(results) / len(content_hashes) if content_hashes else 0,
+        )
+
+        return results
+
+    except (sqlite3.Error, json.JSONDecodeError) as e:
+        msg = f"Failed to retrieve cached embeddings batch: {e}"
+        raise StorageError(
+            code=ErrorCode.STORAGE_SQLITE_QUERY_FAILED,
+            message=msg,
+            context=storage_context(
+                table="embedding_cache",
+                query="get_cached_embeddings_batch",
+            ),
+        ) from e
+
+
+def put_cached_embeddings_batch(
+    conn: Connection,
+    embeddings: list[tuple[str, list[float]]],
+    model_name: str,
+) -> None:
+    """Store multiple embeddings in cache efficiently.
+
+    Uses executemany for batch insertion.
+
+    Args:
+        conn: SQLite connection.
+        embeddings: List of (content_hash, embedding) tuples.
+        model_name: Embedding model identifier.
+
+    Raises:
+        StorageError: If batch cache write fails.
+    """
+    if not embeddings:
+        return
+
+    try:
+        import time
+
+        cursor = conn.cursor()
+        now_ms = int(time.time() * 1000)
+
+        # Prepare batch data
+        batch_data = [
+            (
+                content_hash,
+                json.dumps(embedding),
+                model_name,
+                now_ms,
+            )
+            for content_hash, embedding in embeddings
+        ]
+
+        # Batch insert
+        cursor.executemany(
+            """
+            INSERT OR REPLACE INTO embedding_cache
+                (content_hash, embedding, model_name, created_at, hit_count)
+            VALUES (?, ?, ?, ?, 0)
+            """,
+            batch_data,
+        )
+
+        conn.commit()
+
+        logger.info(
+            "Batch cache write",
+            count=len(embeddings),
+            model=model_name,
+        )
+
+    except (sqlite3.Error, TypeError) as e:
+        msg = f"Failed to cache embeddings batch: {e}"
+        raise StorageError(
+            code=ErrorCode.STORAGE_SQLITE_QUERY_FAILED,
+            message=msg,
+            context=storage_context(
+                table="embedding_cache",
+                query="put_cached_embeddings_batch",
+            ),
+        ) from e
+
+
+def get_cache_stats(conn: Connection) -> dict[str, int]:
+    """Get cache statistics for monitoring.
+
+    Returns:
+        Dict with cache size, total hits, and model breakdown.
+
+    Example:
+        >>> stats = get_cache_stats(conn)
+        >>> stats
+        {
+            'total_entries': 10000,
+            'total_hits': 45000,
+            'models': {'jinaai/jina-embeddings-v2-base-code': 10000}
+        }
+    """
+    try:
+        cursor = conn.cursor()
+
+        # Total entries
+        cursor.execute("SELECT COUNT(*) FROM embedding_cache")
+        total_entries = cursor.fetchone()[0]
+
+        # Total hits
+        cursor.execute("SELECT COALESCE(SUM(hit_count), 0) FROM embedding_cache")
+        total_hits = cursor.fetchone()[0]
+
+        # Model breakdown
+        cursor.execute(
+            """
+            SELECT model_name, COUNT(*)
+            FROM embedding_cache
+            GROUP BY model_name
+            """
+        )
+        models = dict(cursor.fetchall())
+
+        return {
+            "total_entries": total_entries,
+            "total_hits": total_hits,
+            "models": models,
+        }
+
+    except sqlite3.Error as e:
+        logger.warning("Failed to get cache stats", error=str(e))
+        return {"total_entries": 0, "total_hits": 0, "models": {}}
+
+
+def clear_cache(conn: Connection, model_name: str | None = None) -> int:
+    """Clear embedding cache.
+
+    Args:
+        conn: SQLite connection.
+        model_name: Clear only entries for this model, or None for all.
+
+    Returns:
+        Number of entries deleted.
+
+    Raises:
+        StorageError: If cache clear fails.
+    """
+    try:
+        cursor = conn.cursor()
+
+        if model_name:
+            cursor.execute(
+                "DELETE FROM embedding_cache WHERE model_name = ?",
+                (model_name,),
+            )
+        else:
+            cursor.execute("DELETE FROM embedding_cache")
+
+        deleted = cursor.rowcount
+        conn.commit()
+
+        logger.info(
+            "Cache cleared",
+            deleted=deleted,
+            model=model_name or "all",
+        )
+
+        return deleted
+
+    except sqlite3.Error as e:
+        msg = f"Failed to clear cache: {e}"
+        raise StorageError(
+            code=ErrorCode.STORAGE_SQLITE_QUERY_FAILED,
+            message=msg,
+            context=storage_context(
+                table="embedding_cache",
+                query="clear_cache",
+            ),
+        ) from e