contextual-engine 0.1.0__py3-none-any.whl

contextual/embedding/embedder.py

@@ -0,0 +1,414 @@
+"""Production embedder using fastembed with jina-embeddings-v2-base-code.
+
+This module provides a thread-safe, cached embedder singleton for code chunks.
+Uses ONNX int8 quantization for efficient CPU inference.
+
+Key features:
+- Model: jina-embeddings-v2-base-code (768-dim)
+- Batching: Up to 64 documents per batch
+- Caching: Content-hash based deduplication
+- Thread-safe: Singleton with lock protection
+- Fast: ~18 docs/sec on modern CPU
+
+Performance characteristics (from TECHNICAL_SPEC):
+- Model size: ~165MB on disk
+- Context window: 8,192 tokens
+- Batch size: 64 @ 512 tokens
+- Throughput: ~18 docs/sec (int8, 128 seq len)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import threading
+from typing import TYPE_CHECKING
+
+from fastembed import TextEmbedding
+
+from contextual.core.errors import ErrorCode, IndexingError, indexing_context
+from contextual.observability.logging import get_logger
+
+if TYPE_CHECKING:
+    pass
+
+
+logger = get_logger(__name__)
+
+# ============================================================================
+# CONSTANTS (from TECHNICAL_SPEC §4)
+# ============================================================================
+
+MODEL_NAME = "jinaai/jina-embeddings-v2-base-code"
+"""jina-embeddings-v2-base-code - optimized for code understanding."""
+
+DOCS_MODEL_NAME = "nomic-ai/nomic-embed-text-v1.5"
+"""nomic-embed-text-v1.5 - optimized for natural language text (docs)."""
+
+EMBEDDING_DIM = 768
+"""Output dimensionality."""
+
+BATCH_SIZE = 64
+"""Maximum documents per batch for optimal throughput."""
+
+MAX_CONTEXT_LENGTH = 8192
+"""Maximum token context window."""
+
+
+# ============================================================================
+# EMBEDDER CLASS
+# ============================================================================
+
+
+class CodeEmbedder:
+    """Thread-safe code embedder using jina-v2-base-code.
+
+    Wraps fastembed's TextEmbedding with:
+    - Singleton pattern for model reuse
+    - Batched inference for efficiency
+    - Content-hash caching integration
+    - Thread-safe operations
+
+    The model is loaded once and reused across all embedding operations.
+    Model weights are cached by fastembed in ~/.cache/fastembed.
+
+    Attributes:
+        model: fastembed TextEmbedding instance
+        _lock: Threading lock for model access
+    """
+
+    def __init__(self) -> None:
+        """Initialize embedder and load model.
+
+        Downloads model on first use (~165MB).
+        Subsequent runs load from cache.
+
+        Raises:
+            IndexingError: If model loading fails.
+        """
+        self._lock = threading.Lock()
+
+        logger.info(
+            "Initializing code embedder",
+            model=MODEL_NAME,
+            dimensions=EMBEDDING_DIM,
+        )
+
+        try:
+            # Initialize fastembed TextEmbedding
+            # This downloads the model on first use
+            self.model = TextEmbedding(
+                model_name=MODEL_NAME,
+                max_length=MAX_CONTEXT_LENGTH,
+            )
+
+            logger.info(
+                "Code embedder initialized successfully",
+                model=MODEL_NAME,
+            )
+
+        except Exception as e:
+            msg = f"Failed to initialize embedder: {e}"
+            raise IndexingError(
+                code=ErrorCode.INDEXING_EMBEDDING_MODEL_LOAD_FAILED,
+                message=msg,
+                context=indexing_context(),
+            ) from e
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for a batch of texts.
+
+        This is the main entry point for embedding generation.
+        Handles batching automatically if input exceeds BATCH_SIZE.
+
+        Args:
+            texts: List of text strings to embed (code chunks, queries, etc.)
+
+        Returns:
+            List of 768-dimensional embedding vectors.
+
+        Raises:
+            IndexingError: If embedding generation fails.
+
+        Example:
+            >>> embedder = get_embedder()
+            >>> chunks = ["def parse():", "class Config:", "import json"]
+            >>> embeddings = embedder.embed(chunks)
+            >>> len(embeddings), len(embeddings[0])
+            (3, 768)
+        """
+        if not texts:
+            return []
+
+        try:
+            # Process in batches if necessary
+            all_embeddings = []
+
+            for i in range(0, len(texts), BATCH_SIZE):
+                batch = texts[i : i + BATCH_SIZE]
+
+                logger.debug(
+                    "Embedding batch",
+                    batch_size=len(batch),
+                    batch_index=i // BATCH_SIZE,
+                )
+
+                # fastembed returns a generator - convert to list
+                with self._lock:
+                    batch_embeddings = list(self.model.embed(batch))
+
+                # Validate dimensions
+                for idx, emb in enumerate(batch_embeddings):
+                    if len(emb) != EMBEDDING_DIM:
+                        msg = (
+                            f"Unexpected embedding dimension: "
+                            f"expected {EMBEDDING_DIM}, got {len(emb)}"
+                        )
+                        raise IndexingError(
+                            code=ErrorCode.INDEXING_EMBEDDING_INFERENCE_FAILED,
+                            message=msg,
+                            context=indexing_context(),
+                        )
+
+                all_embeddings.extend(batch_embeddings)
+
+            logger.info(
+                "Embeddings generated",
+                count=len(texts),
+                batches=len(range(0, len(texts), BATCH_SIZE)),
+            )
+
+            return [list(emb) for emb in all_embeddings]
+
+        except IndexingError:
+            # Re-raise our own errors
+            raise
+        except Exception as e:
+            msg = f"Embedding inference failed: {e}"
+            raise IndexingError(
+                code=ErrorCode.INDEXING_EMBEDDING_INFERENCE_FAILED,
+                message=msg,
+                context=indexing_context(),
+            ) from e
+
+    def embed_query(self, query: str) -> list[float]:
+        """Embed a single query string.
+
+        Convenience method for embedding search queries.
+        jina-v2-code is symmetric, so no special query prefix needed.
+
+        Args:
+            query: Query string to embed.
+
+        Returns:
+            768-dimensional embedding vector.
+
+        Example:
+            >>> embedder = get_embedder()
+            >>> query_vec = embedder.embed_query("parse JSON configuration")
+            >>> len(query_vec)
+            768
+        """
+        return self.embed([query])[0]
+
+    def compute_content_hash(self, text: str) -> str:
+        """Compute SHA-256 hash of text for caching.
+
+        Used to deduplicate embeddings - identical content gets same hash.
+
+        Args:
+            text: Text content to hash.
+
+        Returns:
+            Hex-encoded SHA-256 hash (64 chars).
+        """
+        return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+class DocsEmbedder:
+    """Thread-safe docs embedder using nomic-embed-text-v1.5.
+
+    Identical to CodeEmbedder in behavior (batching, locking, dimension checks).
+    Only difference is the underlying model name.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+
+        logger.info(
+            "Initializing docs embedder",
+            model=DOCS_MODEL_NAME,
+            dimensions=EMBEDDING_DIM,
+        )
+
+        try:
+            self.model = TextEmbedding(
+                model_name=DOCS_MODEL_NAME,
+                max_length=MAX_CONTEXT_LENGTH,
+            )
+
+            logger.info(
+                "Docs embedder initialized successfully",
+                model=DOCS_MODEL_NAME,
+            )
+
+        except Exception as e:
+            msg = f"Failed to initialize embedder: {e}"
+            raise IndexingError(
+                code=ErrorCode.INDEXING_EMBEDDING_MODEL_LOAD_FAILED,
+                message=msg,
+                context=indexing_context(),
+            ) from e
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+
+        try:
+            all_embeddings = []
+
+            for i in range(0, len(texts), BATCH_SIZE):
+                batch = texts[i : i + BATCH_SIZE]
+
+                logger.debug(
+                    "Embedding batch",
+                    batch_size=len(batch),
+                    batch_index=i // BATCH_SIZE,
+                )
+
+                with self._lock:
+                    batch_embeddings = list(self.model.embed(batch))
+
+                for emb in batch_embeddings:
+                    if len(emb) != EMBEDDING_DIM:
+                        msg = (
+                            f"Unexpected embedding dimension: "
+                            f"expected {EMBEDDING_DIM}, got {len(emb)}"
+                        )
+                        raise IndexingError(
+                            code=ErrorCode.INDEXING_EMBEDDING_INFERENCE_FAILED,
+                            message=msg,
+                            context=indexing_context(),
+                        )
+
+                all_embeddings.extend(batch_embeddings)
+
+            logger.info(
+                "Embeddings generated",
+                count=len(texts),
+                batches=len(range(0, len(texts), BATCH_SIZE)),
+            )
+
+            return [list(emb) for emb in all_embeddings]
+
+        except IndexingError:
+            raise
+        except Exception as e:
+            msg = f"Embedding inference failed: {e}"
+            raise IndexingError(
+                code=ErrorCode.INDEXING_EMBEDDING_INFERENCE_FAILED,
+                message=msg,
+                context=indexing_context(),
+            ) from e
+
+    def embed_query(self, query: str) -> list[float]:
+        return self.embed([query])[0]
+
+    def compute_content_hash(self, text: str) -> str:
+        return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+# ============================================================================
+# SINGLETON PATTERN
+# ============================================================================
+
+_embedder: CodeEmbedder | None = None
+_embedder_lock = threading.Lock()
+
+_docs_embedder: DocsEmbedder | None = None
+_docs_embedder_lock = threading.Lock()
+
+
+def get_embedder() -> CodeEmbedder:
+    """Get or create the global embedder singleton.
+
+    Thread-safe lazy initialization. The embedder is created once
+    on first call and reused for all subsequent calls.
+
+    Returns:
+        Shared CodeEmbedder instance.
+
+    Example:
+        >>> embedder1 = get_embedder()
+        >>> embedder2 = get_embedder()
+        >>> embedder1 is embedder2
+        True
+    """
+    global _embedder
+
+    if _embedder is None:
+        with _embedder_lock:
+            # Double-check pattern
+            if _embedder is None:
+                _embedder = CodeEmbedder()
+
+    return _embedder
+
+
+def get_docs_embedder() -> DocsEmbedder:
+    """Get or create the global docs embedder singleton."""
+    global _docs_embedder
+
+    if _docs_embedder is None:
+        with _docs_embedder_lock:
+            if _docs_embedder is None:
+                _docs_embedder = DocsEmbedder()
+
+    return _docs_embedder
+
+
+def reset_embedder() -> None:
+    """Reset the global embedder singleton.
+
+    Used for testing or to force model reload.
+    Should NOT be called in production code.
+    """
+    global _embedder
+    with _embedder_lock:
+        _embedder = None
+    logger.warning("Embedder singleton reset")
+
+
+# ============================================================================
+# BATCH EMBEDDING HELPERS
+# ============================================================================
+
+
+def embed_with_hashes(
+    texts: list[str],
+) -> list[tuple[str, list[float]]]:
+    """Embed texts and return (content_hash, embedding) pairs.
+
+    Useful for populating embedding cache or detecting duplicates.
+
+    Args:
+        texts: Texts to embed.
+
+    Returns:
+        List of (content_hash, embedding) tuples.
+
+    Example:
+        >>> embedder = get_embedder()
+        >>> results = embed_with_hashes(["def foo():", "class Bar:"])
+        >>> hash1, emb1 = results[0]
+        >>> len(hash1), len(emb1)
+        (64, 768)
+    """
+    embedder = get_embedder()
+
+    # Generate embeddings
+    embeddings = embedder.embed(texts)
+
+    # Compute hashes
+    hashes = [embedder.compute_content_hash(text) for text in texts]
+
+    return list(zip(hashes, embeddings, strict=False))

contextual/embedding/helpers.py

@@ -0,0 +1,252 @@
+"""High-level embedding API with caching integration.
+
+Provides convenience functions that combine embedder + cache for common workflows:
+- embed_chunks_with_cache: Batch embed with cache lookup/write
+- embed_query: Simple query embedding (no cache)
+- warm_cache: Pre-populate cache from database
+
+These are the recommended entry points for indexing and search operations.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from contextual.embedding.cache import (
+    get_cached_embeddings_batch,
+    put_cached_embeddings_batch,
+)
+from contextual.embedding.embedder import MODEL_NAME, get_embedder
+from contextual.observability.logging import get_logger
+
+if TYPE_CHECKING:
+    from sqlite3 import Connection
+
+
+logger = get_logger(__name__)
+
+
+# ============================================================================
+# CACHED EMBEDDING WORKFLOW
+# ============================================================================
+
+
+def embed_chunks(
+    conn: Connection,
+    chunks: list[tuple[str, str]],
+) -> list[tuple[str, list[float]]]:
+    """Embed chunks with cache lookup and write.
+
+    This is the main entry point for indexing pipelines.
+
+    Workflow:
+    1. Extract content hashes from chunks
+    2. Lookup cached embeddings in batch
+    3. Generate embeddings only for cache misses
+    4. Write new embeddings to cache
+    5. Return all embeddings in original order
+
+    Args:
+        conn: SQLite connection (for cache operations).
+        chunks: List of (content_hash, text) tuples.
+
+    Returns:
+        List of (content_hash, embedding) tuples in original order.
+
+    Example:
+        >>> chunks = [
+        ...     ("abc123", "def parse_json():"),
+        ...     ("def456", "class Config:"),
+        ... ]
+        >>> results = embed_chunks(conn, chunks)
+        >>> len(results)
+        2
+        >>> hash1, emb1 = results[0]
+        >>> len(emb1)
+        768
+    """
+    if not chunks:
+        return []
+
+    embedder = get_embedder()
+
+    # Extract hashes and texts
+    hashes = [h for h, _ in chunks]
+    texts = [t for _, t in chunks]
+
+    # Batch cache lookup
+    cached = get_cached_embeddings_batch(conn, hashes, MODEL_NAME)
+
+    # Identify cache misses
+    miss_indices = [i for i, h in enumerate(hashes) if h not in cached]
+    miss_texts = [texts[i] for i in miss_indices]
+    miss_hashes = [hashes[i] for i in miss_indices]
+
+    logger.info(
+        "Embedding with cache",
+        total=len(chunks),
+        cached=len(cached),
+        misses=len(miss_indices),
+        hit_rate=len(cached) / len(chunks) if chunks else 0,
+    )
+
+    # Generate embeddings for misses
+    new_embeddings = {}
+    if miss_texts:
+        logger.debug("Generating embeddings for cache misses", count=len(miss_texts))
+
+        fresh_embeddings = embedder.embed(miss_texts)
+
+        # Build miss hash -> embedding map
+        for hash_val, embedding in zip(miss_hashes, fresh_embeddings, strict=False):
+            new_embeddings[hash_val] = embedding
+
+        # Write to cache
+        cache_entries = list(zip(miss_hashes, fresh_embeddings, strict=False))
+        put_cached_embeddings_batch(conn, cache_entries, MODEL_NAME)
+
+    # Combine cached + new in original order
+    results = []
+    for hash_val in hashes:
+        if hash_val in cached:
+            embedding = cached[hash_val]
+        else:
+            embedding = new_embeddings[hash_val]
+
+        results.append((hash_val, embedding))
+
+    return results
+
+
+def embed_query_simple(query: str) -> list[float]:
+    """Embed a search query (no caching).
+
+    Queries are not cached because:
+    1. They are unique and rarely repeated
+    2. Cache lookup overhead > embedding time for single item
+    3. Query embedding is already fast (~55ms)
+
+    Args:
+        query: Query string.
+
+    Returns:
+        768-dim embedding vector.
+
+    Example:
+        >>> embedding = embed_query_simple("parse JSON configuration")
+        >>> len(embedding)
+        768
+    """
+    embedder = get_embedder()
+    return embedder.embed_query(query)
+
+
+# ============================================================================
+# CACHE WARMING
+# ============================================================================
+
+
+def warm_cache_from_chunks(
+    conn: Connection,
+    chunk_ids: list[int] | None = None,
+) -> int:
+    """Pre-populate embedding cache from existing chunks.
+
+    Useful after:
+    - Fresh database creation
+    - Cache cleared
+    - Model upgrade
+
+    Args:
+        conn: SQLite connection.
+        chunk_ids: Specific chunk IDs to warm, or None for all.
+
+    Returns:
+        Number of embeddings cached.
+
+    Example:
+        >>> count = warm_cache_from_chunks(conn)
+        >>> print(f"Cached {count} embeddings")
+    """
+    try:
+        cursor = conn.cursor()
+
+        # Fetch chunks
+        if chunk_ids:
+            placeholders = ",".join("?" * len(chunk_ids))
+            query = f"""
+                SELECT content_hash, content
+                FROM chunks
+                WHERE id IN ({placeholders})
+                AND content_hash IS NOT NULL
+            """
+            cursor.execute(query, chunk_ids)
+        else:
+            query = """
+                SELECT content_hash, content
+                FROM chunks
+                WHERE content_hash IS NOT NULL
+            """
+            cursor.execute(query)
+
+        rows = cursor.fetchall()
+        chunks = [(h, c) for h, c in rows]
+
+        logger.info("Warming cache from chunks", chunk_count=len(chunks))
+
+        # Use cached embedding workflow
+        # (will skip already-cached entries automatically)
+        results = embed_chunks(conn, chunks)
+
+        return len(results)
+
+    except Exception as e:
+        logger.error("Cache warming failed", error=str(e))
+        return 0
+
+
+# ============================================================================
+# STATISTICS
+# ============================================================================
+
+
+def get_embedding_stats(conn: Connection) -> dict[str, int | float]:
+    """Get embedding and cache statistics.
+
+    Returns:
+        Dict with cache stats and efficiency metrics.
+
+    Example:
+        >>> stats = get_embedding_stats(conn)
+        >>> stats
+        {
+            'cache_entries': 10000,
+            'cache_hits': 45000,
+            'cache_hit_rate': 0.82,
+            'estimated_time_saved_hours': 12.5
+        }
+    """
+    from contextual.embedding.cache import get_cache_stats
+
+    cache_stats = get_cache_stats(conn)
+
+    # Calculate efficiency metrics
+    total_entries = cache_stats["total_entries"]
+    total_hits = cache_stats["total_hits"]
+
+    # Assume ~55ms per embedding generation
+    # Cache hit saves ~55ms, lookup costs ~0.1ms
+    time_saved_ms = total_hits * 55
+    time_saved_hours = time_saved_ms / 1000 / 60 / 60
+
+    cache_hit_rate = 0.0
+    if total_entries + total_hits > 0:
+        cache_hit_rate = total_hits / (total_entries + total_hits)
+
+    return {
+        "cache_entries": total_entries,
+        "cache_hits": total_hits,
+        "cache_hit_rate": round(cache_hit_rate, 3),
+        "estimated_time_saved_hours": round(time_saved_hours, 2),
+        "models": cache_stats["models"],
+    }

contextual/git/__init__.py

@@ -0,0 +1,22 @@
+"""Git integration module for temporal code attribution.
+
+Public API:
+- GitRepository: Wrapper for git operations
+- BlameInfo: Line range attribution
+- ChunkAttribution: Aggregated chunk metadata
+"""
+
+from __future__ import annotations
+
+from contextual.git.blame import (
+    BlameInfo,
+    ChunkAttribution,
+    GitRepository,
+)
+
+
+__all__ = [
+    "BlameInfo",
+    "ChunkAttribution",
+    "GitRepository",
+]