keep-skill 0.1.0__py3-none-any.whl

keep/providers/embedding_cache.py

@@ -0,0 +1,260 @@
+"""
+Embedding cache using SQLite.
+
+Wraps any EmbeddingProvider to cache embeddings by content hash,
+avoiding redundant embedding calls for unchanged content.
+"""
+
+import hashlib
+import json
+import sqlite3
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+from .base import EmbeddingProvider
+
+
+class EmbeddingCache:
+    """
+    SQLite-based embedding cache.
+    
+    Cache key is SHA256(model_name + content), so different models
+    don't share cached embeddings.
+    """
+    
+    def __init__(self, cache_path: Path, max_entries: int = 50000):
+        """
+        Args:
+            cache_path: Path to SQLite database file
+            max_entries: Maximum cache entries (LRU eviction when exceeded)
+        """
+        self._cache_path = cache_path
+        self._max_entries = max_entries
+        self._conn: Optional[sqlite3.Connection] = None
+        self._init_db()
+    
+    def _init_db(self) -> None:
+        """Initialize the SQLite database."""
+        self._cache_path.parent.mkdir(parents=True, exist_ok=True)
+        self._conn = sqlite3.connect(str(self._cache_path), check_same_thread=False)
+        self._conn.execute("""
+            CREATE TABLE IF NOT EXISTS embedding_cache (
+                content_hash TEXT PRIMARY KEY,
+                model_name TEXT NOT NULL,
+                embedding BLOB NOT NULL,
+                dimension INTEGER NOT NULL,
+                created_at TEXT NOT NULL,
+                last_accessed TEXT NOT NULL
+            )
+        """)
+        self._conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_last_accessed 
+            ON embedding_cache(last_accessed)
+        """)
+        self._conn.commit()
+    
+    def _hash_key(self, model_name: str, content: str) -> str:
+        """Generate cache key from model name and content."""
+        key_input = f"{model_name}:{content}"
+        return hashlib.sha256(key_input.encode("utf-8")).hexdigest()
+    
+    def get(self, model_name: str, content: str) -> Optional[list[float]]:
+        """
+        Get cached embedding if it exists.
+        
+        Updates last_accessed timestamp on hit.
+        """
+        content_hash = self._hash_key(model_name, content)
+        cursor = self._conn.execute(
+            "SELECT embedding FROM embedding_cache WHERE content_hash = ?",
+            (content_hash,)
+        )
+        row = cursor.fetchone()
+        
+        if row is not None:
+            # Update last_accessed
+            now = datetime.now(timezone.utc).isoformat()
+            self._conn.execute(
+                "UPDATE embedding_cache SET last_accessed = ? WHERE content_hash = ?",
+                (now, content_hash)
+            )
+            self._conn.commit()
+            
+            # Deserialize embedding
+            return json.loads(row[0])
+        
+        return None
+    
+    def put(
+        self,
+        model_name: str,
+        content: str,
+        embedding: list[float]
+    ) -> None:
+        """
+        Cache an embedding.
+        
+        Evicts oldest entries if cache exceeds max_entries.
+        """
+        content_hash = self._hash_key(model_name, content)
+        now = datetime.now(timezone.utc).isoformat()
+        embedding_blob = json.dumps(embedding)
+        
+        self._conn.execute("""
+            INSERT OR REPLACE INTO embedding_cache 
+            (content_hash, model_name, embedding, dimension, created_at, last_accessed)
+            VALUES (?, ?, ?, ?, ?, ?)
+        """, (content_hash, model_name, embedding_blob, len(embedding), now, now))
+        self._conn.commit()
+        
+        # Evict old entries if needed
+        self._maybe_evict()
+    
+    def _maybe_evict(self) -> None:
+        """Evict oldest entries if cache exceeds max size."""
+        cursor = self._conn.execute("SELECT COUNT(*) FROM embedding_cache")
+        count = cursor.fetchone()[0]
+        
+        if count > self._max_entries:
+            # Delete oldest 10% by last_accessed
+            evict_count = max(1, count // 10)
+            self._conn.execute("""
+                DELETE FROM embedding_cache 
+                WHERE content_hash IN (
+                    SELECT content_hash FROM embedding_cache 
+                    ORDER BY last_accessed ASC 
+                    LIMIT ?
+                )
+            """, (evict_count,))
+            self._conn.commit()
+    
+    def stats(self) -> dict:
+        """Get cache statistics."""
+        cursor = self._conn.execute("""
+            SELECT COUNT(*), COUNT(DISTINCT model_name) 
+            FROM embedding_cache
+        """)
+        count, models = cursor.fetchone()
+        return {
+            "entries": count,
+            "models": models,
+            "max_entries": self._max_entries,
+            "cache_path": str(self._cache_path),
+        }
+    
+    def clear(self) -> None:
+        """Clear all cached embeddings."""
+        self._conn.execute("DELETE FROM embedding_cache")
+        self._conn.commit()
+
+    def close(self) -> None:
+        """Close the database connection."""
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+
+    def __del__(self) -> None:
+        """Ensure connection is closed on cleanup."""
+        self.close()
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - close connection."""
+        self.close()
+        return False
+
+
+class CachingEmbeddingProvider:
+    """
+    Wrapper that adds caching to any EmbeddingProvider.
+    
+    Usage:
+        base_provider = SentenceTransformerEmbedding()
+        cached = CachingEmbeddingProvider(base_provider, cache_path)
+    """
+    
+    def __init__(
+        self,
+        provider: EmbeddingProvider,
+        cache_path: Path,
+        max_entries: int = 50000
+    ):
+        self._provider = provider
+        self._cache = EmbeddingCache(cache_path, max_entries)
+        self._hits = 0
+        self._misses = 0
+    
+    @property
+    def model_name(self) -> str:
+        """Get the underlying provider's model name."""
+        return getattr(self._provider, "model_name", "unknown")
+    
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension from the wrapped provider."""
+        return self._provider.dimension
+    
+    def embed(self, text: str) -> list[float]:
+        """
+        Get embedding, using cache when available.
+        """
+        # Check cache
+        cached = self._cache.get(self.model_name, text)
+        if cached is not None:
+            self._hits += 1
+            return cached
+        
+        # Cache miss - compute embedding
+        self._misses += 1
+        embedding = self._provider.embed(text)
+        
+        # Store in cache
+        self._cache.put(self.model_name, text, embedding)
+        
+        return embedding
+    
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """
+        Get embeddings for batch, using cache where available.
+        
+        Only computes embeddings for cache misses.
+        """
+        results: list[Optional[list[float]]] = [None] * len(texts)
+        to_embed: list[tuple[int, str]] = []
+        
+        # Check cache for each text
+        for i, text in enumerate(texts):
+            cached = self._cache.get(self.model_name, text)
+            if cached is not None:
+                self._hits += 1
+                results[i] = cached
+            else:
+                self._misses += 1
+                to_embed.append((i, text))
+        
+        # Batch embed cache misses
+        if to_embed:
+            indices, texts_to_embed = zip(*to_embed)
+            embeddings = self._provider.embed_batch(list(texts_to_embed))
+            
+            for idx, text, embedding in zip(indices, texts_to_embed, embeddings):
+                results[idx] = embedding
+                self._cache.put(self.model_name, text, embedding)
+        
+        return results  # type: ignore
+    
+    def stats(self) -> dict:
+        """Get cache and hit/miss statistics."""
+        cache_stats = self._cache.stats()
+        total = self._hits + self._misses
+        hit_rate = self._hits / total if total > 0 else 0.0
+        return {
+            **cache_stats,
+            "hits": self._hits,
+            "misses": self._misses,
+            "hit_rate": f"{hit_rate:.1%}",
+        }

keep/providers/embeddings.py

@@ -0,0 +1,245 @@
+"""
+Embedding providers for generating vector representations of text.
+"""
+
+import os
+from typing import Any
+
+from .base import EmbeddingProvider, get_registry
+
+
+class SentenceTransformerEmbedding:
+    """
+    Embedding provider using sentence-transformers library.
+    
+    Runs locally, no API key required. Good default for getting started.
+    
+    Requires: pip install sentence-transformers
+    """
+    
+    def __init__(self, model: str = "all-MiniLM-L6-v2"):
+        """
+        Args:
+            model: Model name from sentence-transformers hub
+        """
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise RuntimeError(
+                "SentenceTransformerEmbedding requires 'sentence-transformers' library. "
+                "Install with: pip install sentence-transformers"
+            )
+        
+        self.model_name = model
+        self._model = SentenceTransformer(model)
+    
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension from the model."""
+        return self._model.get_sentence_embedding_dimension()
+    
+    def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text."""
+        embedding = self._model.encode(text, convert_to_numpy=True)
+        return embedding.tolist()
+    
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts."""
+        embeddings = self._model.encode(texts, convert_to_numpy=True)
+        return embeddings.tolist()
+
+
+class OpenAIEmbedding:
+    """
+    Embedding provider using OpenAI's API.
+    
+    Requires: KEEP_OPENAI_API_KEY or OPENAI_API_KEY environment variable.
+    Requires: pip install openai
+    """
+    
+    # Model dimensions (as of 2024)
+    MODEL_DIMENSIONS = {
+        "text-embedding-3-small": 1536,
+        "text-embedding-3-large": 3072,
+        "text-embedding-ada-002": 1536,
+    }
+    
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        api_key: str | None = None,
+    ):
+        """
+        Args:
+            model: OpenAI embedding model name
+            api_key: API key (defaults to environment variable)
+        """
+        try:
+            from openai import OpenAI
+        except ImportError:
+            raise RuntimeError(
+                "OpenAIEmbedding requires 'openai' library. "
+                "Install with: pip install openai"
+            )
+        
+        self.model_name = model
+        self._dimension = self.MODEL_DIMENSIONS.get(model, 1536)
+        
+        # Resolve API key
+        key = api_key or os.environ.get("KEEP_OPENAI_API_KEY") or os.environ.get("OPENAI_API_KEY")
+        if not key:
+            raise ValueError(
+                "OpenAI API key required. Set KEEP_OPENAI_API_KEY or OPENAI_API_KEY"
+            )
+        
+        self._client = OpenAI(api_key=key)
+    
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension for the model."""
+        return self._dimension
+    
+    def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text."""
+        response = self._client.embeddings.create(
+            model=self.model_name,
+            input=text,
+        )
+        return response.data[0].embedding
+    
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts."""
+        response = self._client.embeddings.create(
+            model=self.model_name,
+            input=texts,
+        )
+        # Sort by index to ensure order matches input
+        sorted_data = sorted(response.data, key=lambda x: x.index)
+        return [d.embedding for d in sorted_data]
+
+
+class GeminiEmbedding:
+    """
+    Embedding provider using Google's Gemini API.
+
+    Requires: GEMINI_API_KEY or GOOGLE_API_KEY environment variable.
+    Requires: pip install google-genai
+    """
+
+    # Model dimensions (as of 2025)
+    MODEL_DIMENSIONS = {
+        "text-embedding-004": 768,
+        "embedding-001": 768,
+        "gemini-embedding-001": 768,
+    }
+
+    def __init__(
+        self,
+        model: str = "text-embedding-004",
+        api_key: str | None = None,
+    ):
+        """
+        Args:
+            model: Gemini embedding model name
+            api_key: API key (defaults to environment variable)
+        """
+        try:
+            from google import genai
+        except ImportError:
+            raise RuntimeError(
+                "GeminiEmbedding requires 'google-genai' library. "
+                "Install with: pip install google-genai"
+            )
+
+        self.model_name = model
+        self._dimension = self.MODEL_DIMENSIONS.get(model, 768)
+
+        # Resolve API key
+        key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+        if not key:
+            raise ValueError(
+                "Gemini API key required. Set GEMINI_API_KEY or GOOGLE_API_KEY"
+            )
+
+        self._client = genai.Client(api_key=key)
+
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension for the model."""
+        return self._dimension
+
+    def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text."""
+        result = self._client.models.embed_content(
+            model=self.model_name,
+            contents=text,
+        )
+        return list(result.embeddings[0].values)
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts."""
+        result = self._client.models.embed_content(
+            model=self.model_name,
+            contents=texts,
+        )
+        return [list(e.values) for e in result.embeddings]
+
+
+class OllamaEmbedding:
+    """
+    Embedding provider using Ollama's local API.
+
+    Requires: Ollama running locally (default: http://localhost:11434)
+    """
+    
+    def __init__(
+        self,
+        model: str = "nomic-embed-text",
+        base_url: str = "http://localhost:11434",
+    ):
+        """
+        Args:
+            model: Ollama model name
+            base_url: Ollama API base URL
+        """
+        self.model_name = model
+        self.base_url = base_url.rstrip("/")
+        self._dimension: int | None = None
+    
+    @property
+    def dimension(self) -> int:
+        """Get embedding dimension (determined on first embed call)."""
+        if self._dimension is None:
+            # Generate a test embedding to determine dimension
+            test_embedding = self.embed("test")
+            self._dimension = len(test_embedding)
+        return self._dimension
+    
+    def embed(self, text: str) -> list[float]:
+        """Generate embedding for a single text."""
+        import requests
+        
+        response = requests.post(
+            f"{self.base_url}/api/embeddings",
+            json={"model": self.model_name, "prompt": text},
+        )
+        response.raise_for_status()
+        
+        embedding = response.json()["embedding"]
+        
+        if self._dimension is None:
+            self._dimension = len(embedding)
+        
+        return embedding
+    
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for multiple texts (sequential for Ollama)."""
+        return [self.embed(text) for text in texts]
+
+
+# Register providers
+_registry = get_registry()
+_registry.register_embedding("sentence-transformers", SentenceTransformerEmbedding)
+_registry.register_embedding("openai", OpenAIEmbedding)
+_registry.register_embedding("gemini", GeminiEmbedding)
+_registry.register_embedding("ollama", OllamaEmbedding)