claude-code-workflow 6.3.18 → 6.3.19

package/codex-lens/src/codexlens/indexing/embedding.py

@@ -0,0 +1,582 @@
+"""Multi-type embedding backends for cascade retrieval.
+
+This module provides embedding backends optimized for cascade retrieval:
+1. BinaryEmbeddingBackend - Fast coarse filtering with binary vectors
+2. DenseEmbeddingBackend - High-precision dense vectors for reranking
+3. CascadeEmbeddingBackend - Combined binary + dense for two-stage retrieval
+
+Cascade retrieval workflow:
+1. Binary search (fast, ~32 bytes/vector) -> top-K candidates
+2. Dense rerank (precise, ~8KB/vector) -> final results
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Iterable, List, Optional, Tuple
+
+import numpy as np
+
+from codexlens.semantic.base import BaseEmbedder
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Utility Functions
+# =============================================================================
+
+
+def binarize_embedding(embedding: np.ndarray) -> np.ndarray:
+    """Convert float embedding to binary vector.
+
+    Applies sign-based quantization: values > 0 become 1, values <= 0 become 0.
+
+    Args:
+        embedding: Float32 embedding of any dimension
+
+    Returns:
+        Binary vector (uint8 with values 0 or 1) of same dimension
+    """
+    return (embedding > 0).astype(np.uint8)
+
+
+def pack_binary_embedding(binary_vector: np.ndarray) -> bytes:
+    """Pack binary vector into compact bytes format.
+
+    Packs 8 binary values into each byte for storage efficiency.
+    For a 256-dim binary vector, output is 32 bytes.
+
+    Args:
+        binary_vector: Binary vector (uint8 with values 0 or 1)
+
+    Returns:
+        Packed bytes (length = ceil(dim / 8))
+    """
+    # Ensure vector length is multiple of 8 by padding if needed
+    dim = len(binary_vector)
+    padded_dim = ((dim + 7) // 8) * 8
+    if padded_dim > dim:
+        padded = np.zeros(padded_dim, dtype=np.uint8)
+        padded[:dim] = binary_vector
+        binary_vector = padded
+
+    # Pack 8 bits per byte
+    packed = np.packbits(binary_vector)
+    return packed.tobytes()
+
+
+def unpack_binary_embedding(packed_bytes: bytes, dim: int = 256) -> np.ndarray:
+    """Unpack bytes back to binary vector.
+
+    Args:
+        packed_bytes: Packed binary data
+        dim: Original vector dimension (default: 256)
+
+    Returns:
+        Binary vector (uint8 with values 0 or 1)
+    """
+    unpacked = np.unpackbits(np.frombuffer(packed_bytes, dtype=np.uint8))
+    return unpacked[:dim]
+
+
+def hamming_distance(a: bytes, b: bytes) -> int:
+    """Compute Hamming distance between two packed binary vectors.
+
+    Uses XOR and popcount for efficient distance computation.
+
+    Args:
+        a: First packed binary vector
+        b: Second packed binary vector
+
+    Returns:
+        Hamming distance (number of differing bits)
+    """
+    a_arr = np.frombuffer(a, dtype=np.uint8)
+    b_arr = np.frombuffer(b, dtype=np.uint8)
+    xor = np.bitwise_xor(a_arr, b_arr)
+    return int(np.unpackbits(xor).sum())
+
+
+# =============================================================================
+# Binary Embedding Backend
+# =============================================================================
+
+
+class BinaryEmbeddingBackend(BaseEmbedder):
+    """Generate 256-dimensional binary embeddings for fast coarse retrieval.
+
+    Uses a lightweight embedding model and applies sign-based quantization
+    to produce compact binary vectors (32 bytes per embedding).
+
+    Suitable for:
+    - First-stage candidate retrieval
+    - Hamming distance-based similarity search
+    - Memory-constrained environments
+
+    Model: sentence-transformers/all-MiniLM-L6-v2 (384 dim) -> quantized to 256 bits
+    """
+
+    DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"  # 384 dim, fast
+    BINARY_DIM = 256
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        use_gpu: bool = True,
+    ) -> None:
+        """Initialize binary embedding backend.
+
+        Args:
+            model_name: Base embedding model name. Defaults to BAAI/bge-small-en-v1.5
+            use_gpu: Whether to use GPU acceleration
+        """
+        from codexlens.semantic import SEMANTIC_AVAILABLE
+
+        if not SEMANTIC_AVAILABLE:
+            raise ImportError(
+                "Semantic search dependencies not available. "
+                "Install with: pip install codexlens[semantic]"
+            )
+
+        self._model_name = model_name or self.DEFAULT_MODEL
+        self._use_gpu = use_gpu
+        self._model = None
+
+        # Projection matrix for dimension reduction (lazily initialized)
+        self._projection_matrix: Optional[np.ndarray] = None
+
+    @property
+    def model_name(self) -> str:
+        """Return model name."""
+        return self._model_name
+
+    @property
+    def embedding_dim(self) -> int:
+        """Return binary embedding dimension (256)."""
+        return self.BINARY_DIM
+
+    @property
+    def packed_bytes(self) -> int:
+        """Return packed bytes size (32 bytes for 256 bits)."""
+        return self.BINARY_DIM // 8
+
+    def _load_model(self) -> None:
+        """Lazy load the embedding model."""
+        if self._model is not None:
+            return
+
+        from fastembed import TextEmbedding
+        from codexlens.semantic.gpu_support import get_optimal_providers
+
+        providers = get_optimal_providers(use_gpu=self._use_gpu, with_device_options=True)
+        try:
+            self._model = TextEmbedding(
+                model_name=self._model_name,
+                providers=providers,
+            )
+        except TypeError:
+            # Fallback for older fastembed versions
+            self._model = TextEmbedding(model_name=self._model_name)
+
+        logger.debug(f"BinaryEmbeddingBackend loaded model: {self._model_name}")
+
+    def _get_projection_matrix(self, input_dim: int) -> np.ndarray:
+        """Get or create projection matrix for dimension reduction.
+
+        Uses random projection with fixed seed for reproducibility.
+
+        Args:
+            input_dim: Input embedding dimension from base model
+
+        Returns:
+            Projection matrix of shape (input_dim, BINARY_DIM)
+        """
+        if self._projection_matrix is not None:
+            return self._projection_matrix
+
+        # Fixed seed for reproducibility across sessions
+        rng = np.random.RandomState(42)
+        # Gaussian random projection
+        self._projection_matrix = rng.randn(input_dim, self.BINARY_DIM).astype(np.float32)
+        # Normalize columns for consistent scale
+        norms = np.linalg.norm(self._projection_matrix, axis=0, keepdims=True)
+        self._projection_matrix /= (norms + 1e-8)
+
+        return self._projection_matrix
+
+    def embed_to_numpy(self, texts: str | Iterable[str]) -> np.ndarray:
+        """Generate binary embeddings as numpy array.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            Binary embeddings of shape (n_texts, 256) with values 0 or 1
+        """
+        self._load_model()
+
+        if isinstance(texts, str):
+            texts = [texts]
+        else:
+            texts = list(texts)
+
+        # Get base float embeddings
+        float_embeddings = np.array(list(self._model.embed(texts)))
+        input_dim = float_embeddings.shape[1]
+
+        # Project to target dimension if needed
+        if input_dim != self.BINARY_DIM:
+            projection = self._get_projection_matrix(input_dim)
+            float_embeddings = float_embeddings @ projection
+
+        # Binarize
+        return binarize_embedding(float_embeddings)
+
+    def embed_packed(self, texts: str | Iterable[str]) -> List[bytes]:
+        """Generate packed binary embeddings.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            List of packed bytes (32 bytes each for 256-dim)
+        """
+        binary = self.embed_to_numpy(texts)
+        return [pack_binary_embedding(vec) for vec in binary]
+
+
+# =============================================================================
+# Dense Embedding Backend
+# =============================================================================
+
+
+class DenseEmbeddingBackend(BaseEmbedder):
+    """Generate high-dimensional dense embeddings for precise reranking.
+
+    Uses large embedding models to produce 2048-dimensional float32 vectors
+    for maximum retrieval quality.
+
+    Suitable for:
+    - Second-stage reranking
+    - High-precision similarity search
+    - Quality-critical applications
+
+    Model: BAAI/bge-large-en-v1.5 (1024 dim) with optional expansion
+    """
+
+    DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"  # 384 dim, use small for testing
+    TARGET_DIM = 768  # Reduced target for faster testing
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        use_gpu: bool = True,
+        expand_dim: bool = True,
+    ) -> None:
+        """Initialize dense embedding backend.
+
+        Args:
+            model_name: Dense embedding model name. Defaults to BAAI/bge-large-en-v1.5
+            use_gpu: Whether to use GPU acceleration
+            expand_dim: If True, expand embeddings to TARGET_DIM using learned expansion
+        """
+        from codexlens.semantic import SEMANTIC_AVAILABLE
+
+        if not SEMANTIC_AVAILABLE:
+            raise ImportError(
+                "Semantic search dependencies not available. "
+                "Install with: pip install codexlens[semantic]"
+            )
+
+        self._model_name = model_name or self.DEFAULT_MODEL
+        self._use_gpu = use_gpu
+        self._expand_dim = expand_dim
+        self._model = None
+        self._native_dim: Optional[int] = None
+
+        # Expansion matrix for dimension expansion (lazily initialized)
+        self._expansion_matrix: Optional[np.ndarray] = None
+
+    @property
+    def model_name(self) -> str:
+        """Return model name."""
+        return self._model_name
+
+    @property
+    def embedding_dim(self) -> int:
+        """Return embedding dimension.
+
+        Returns TARGET_DIM if expand_dim is True, otherwise native model dimension.
+        """
+        if self._expand_dim:
+            return self.TARGET_DIM
+        # Return cached native dim or estimate based on model
+        if self._native_dim is not None:
+            return self._native_dim
+        # Model dimension estimates
+        model_dims = {
+            "BAAI/bge-large-en-v1.5": 1024,
+            "BAAI/bge-base-en-v1.5": 768,
+            "BAAI/bge-small-en-v1.5": 384,
+            "intfloat/multilingual-e5-large": 1024,
+        }
+        return model_dims.get(self._model_name, 1024)
+
+    @property
+    def max_tokens(self) -> int:
+        """Return maximum token limit."""
+        return 512  # Conservative default for large models
+
+    def _load_model(self) -> None:
+        """Lazy load the embedding model."""
+        if self._model is not None:
+            return
+
+        from fastembed import TextEmbedding
+        from codexlens.semantic.gpu_support import get_optimal_providers
+
+        providers = get_optimal_providers(use_gpu=self._use_gpu, with_device_options=True)
+        try:
+            self._model = TextEmbedding(
+                model_name=self._model_name,
+                providers=providers,
+            )
+        except TypeError:
+            self._model = TextEmbedding(model_name=self._model_name)
+
+        logger.debug(f"DenseEmbeddingBackend loaded model: {self._model_name}")
+
+    def _get_expansion_matrix(self, input_dim: int) -> np.ndarray:
+        """Get or create expansion matrix for dimension expansion.
+
+        Uses random orthogonal projection for information-preserving expansion.
+
+        Args:
+            input_dim: Input embedding dimension from base model
+
+        Returns:
+            Expansion matrix of shape (input_dim, TARGET_DIM)
+        """
+        if self._expansion_matrix is not None:
+            return self._expansion_matrix
+
+        # Fixed seed for reproducibility
+        rng = np.random.RandomState(123)
+
+        # Create semi-orthogonal expansion matrix
+        # First input_dim columns form identity-like structure
+        self._expansion_matrix = np.zeros((input_dim, self.TARGET_DIM), dtype=np.float32)
+
+        # Copy original dimensions
+        copy_dim = min(input_dim, self.TARGET_DIM)
+        self._expansion_matrix[:copy_dim, :copy_dim] = np.eye(copy_dim, dtype=np.float32)
+
+        # Fill remaining with random projections
+        if self.TARGET_DIM > input_dim:
+            random_part = rng.randn(input_dim, self.TARGET_DIM - input_dim).astype(np.float32)
+            # Normalize
+            norms = np.linalg.norm(random_part, axis=0, keepdims=True)
+            random_part /= (norms + 1e-8)
+            self._expansion_matrix[:, input_dim:] = random_part
+
+        return self._expansion_matrix
+
+    def embed_to_numpy(self, texts: str | Iterable[str]) -> np.ndarray:
+        """Generate dense embeddings as numpy array.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            Dense embeddings of shape (n_texts, TARGET_DIM) as float32
+        """
+        self._load_model()
+
+        if isinstance(texts, str):
+            texts = [texts]
+        else:
+            texts = list(texts)
+
+        # Get base float embeddings
+        float_embeddings = np.array(list(self._model.embed(texts)), dtype=np.float32)
+        self._native_dim = float_embeddings.shape[1]
+
+        # Expand to target dimension if needed
+        if self._expand_dim and self._native_dim < self.TARGET_DIM:
+            expansion = self._get_expansion_matrix(self._native_dim)
+            float_embeddings = float_embeddings @ expansion
+
+        return float_embeddings
+
+
+# =============================================================================
+# Cascade Embedding Backend
+# =============================================================================
+
+
+class CascadeEmbeddingBackend(BaseEmbedder):
+    """Combined binary + dense embedding backend for cascade retrieval.
+
+    Generates both binary (for fast coarse filtering) and dense (for precise
+    reranking) embeddings in a single pass, optimized for two-stage retrieval.
+
+    Cascade workflow:
+    1. encode_cascade() returns (binary_embeddings, dense_embeddings)
+    2. Binary search: Use Hamming distance on binary vectors -> top-K candidates
+    3. Dense rerank: Use cosine similarity on dense vectors -> final results
+
+    Memory efficiency:
+    - Binary: 32 bytes per vector (256 bits)
+    - Dense: 8192 bytes per vector (2048 x float32)
+    - Total: ~8KB per document for full cascade support
+    """
+
+    def __init__(
+        self,
+        binary_model: Optional[str] = None,
+        dense_model: Optional[str] = None,
+        use_gpu: bool = True,
+    ) -> None:
+        """Initialize cascade embedding backend.
+
+        Args:
+            binary_model: Model for binary embeddings. Defaults to BAAI/bge-small-en-v1.5
+            dense_model: Model for dense embeddings. Defaults to BAAI/bge-large-en-v1.5
+            use_gpu: Whether to use GPU acceleration
+        """
+        self._binary_backend = BinaryEmbeddingBackend(
+            model_name=binary_model,
+            use_gpu=use_gpu,
+        )
+        self._dense_backend = DenseEmbeddingBackend(
+            model_name=dense_model,
+            use_gpu=use_gpu,
+            expand_dim=True,
+        )
+        self._use_gpu = use_gpu
+
+    @property
+    def model_name(self) -> str:
+        """Return model names for both backends."""
+        return f"cascade({self._binary_backend.model_name}, {self._dense_backend.model_name})"
+
+    @property
+    def embedding_dim(self) -> int:
+        """Return dense embedding dimension (for compatibility)."""
+        return self._dense_backend.embedding_dim
+
+    @property
+    def binary_dim(self) -> int:
+        """Return binary embedding dimension."""
+        return self._binary_backend.embedding_dim
+
+    @property
+    def dense_dim(self) -> int:
+        """Return dense embedding dimension."""
+        return self._dense_backend.embedding_dim
+
+    def embed_to_numpy(self, texts: str | Iterable[str]) -> np.ndarray:
+        """Generate dense embeddings (for BaseEmbedder compatibility).
+
+        For cascade embeddings, use encode_cascade() instead.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            Dense embeddings of shape (n_texts, dense_dim)
+        """
+        return self._dense_backend.embed_to_numpy(texts)
+
+    def encode_cascade(
+        self,
+        texts: str | Iterable[str],
+        batch_size: int = 32,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """Generate both binary and dense embeddings.
+
+        Args:
+            texts: Single text or iterable of texts
+            batch_size: Batch size for processing
+
+        Returns:
+            Tuple of:
+            - binary_embeddings: Shape (n_texts, 256), uint8 values 0/1
+            - dense_embeddings: Shape (n_texts, 2048), float32
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        else:
+            texts = list(texts)
+
+        binary_embeddings = self._binary_backend.embed_to_numpy(texts)
+        dense_embeddings = self._dense_backend.embed_to_numpy(texts)
+
+        return binary_embeddings, dense_embeddings
+
+    def encode_binary(self, texts: str | Iterable[str]) -> np.ndarray:
+        """Generate only binary embeddings.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            Binary embeddings of shape (n_texts, 256)
+        """
+        return self._binary_backend.embed_to_numpy(texts)
+
+    def encode_dense(self, texts: str | Iterable[str]) -> np.ndarray:
+        """Generate only dense embeddings.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            Dense embeddings of shape (n_texts, 2048)
+        """
+        return self._dense_backend.embed_to_numpy(texts)
+
+    def encode_binary_packed(self, texts: str | Iterable[str]) -> List[bytes]:
+        """Generate packed binary embeddings.
+
+        Args:
+            texts: Single text or iterable of texts
+
+        Returns:
+            List of packed bytes (32 bytes each)
+        """
+        return self._binary_backend.embed_packed(texts)
+
+
+# =============================================================================
+# Factory Function
+# =============================================================================
+
+
+def get_cascade_embedder(
+    binary_model: Optional[str] = None,
+    dense_model: Optional[str] = None,
+    use_gpu: bool = True,
+) -> CascadeEmbeddingBackend:
+    """Factory function to create a cascade embedder.
+
+    Args:
+        binary_model: Model for binary embeddings (default: BAAI/bge-small-en-v1.5)
+        dense_model: Model for dense embeddings (default: BAAI/bge-large-en-v1.5)
+        use_gpu: Whether to use GPU acceleration
+
+    Returns:
+        Configured CascadeEmbeddingBackend instance
+
+    Example:
+        >>> embedder = get_cascade_embedder()
+        >>> binary, dense = embedder.encode_cascade(["hello world"])
+        >>> binary.shape  # (1, 256)
+        >>> dense.shape   # (1, 2048)
+    """
+    return CascadeEmbeddingBackend(
+        binary_model=binary_model,
+        dense_model=dense_model,
+        use_gpu=use_gpu,
+    )

package/codex-lens/src/codexlens/indexing/symbol_extractor.py

@@ -4,6 +4,11 @@ import sqlite3
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
+try:
+    from codexlens.parsers.treesitter_parser import TreeSitterSymbolParser
+except Exception:  # pragma: no cover - optional dependency / platform variance
+    TreeSitterSymbolParser = None  # type: ignore[assignment]
+
 
 class SymbolExtractor:
     """Extract symbols and relationships from source code using regex patterns."""
@@ -118,7 +123,7 @@ class SymbolExtractor:
 
         patterns = self.PATTERNS[lang]
         symbols = []
-        relationships = []
+        relationships: List[Dict] = []
         lines = content.split('\n')
 
         current_scope = None
@@ -141,33 +146,62 @@ class SymbolExtractor:
                         })
                         current_scope = name
 
-            # Extract imports
-            if 'import' in patterns:
-                match = re.search(patterns['import'], line)
-                if match:
-                    import_target = match.group(1) or match.group(2) if match.lastindex >= 2 else match.group(1)
-                    if import_target and current_scope:
-                        relationships.append({
-                            'source_scope': current_scope,
-                            'target': import_target.strip(),
-                            'type': 'imports',
-                            'file_path': str(file_path),
-                            'line': line_num,
-                        })
-
-            # Extract function calls (simplified)
-            if 'call' in patterns and current_scope:
-                for match in re.finditer(patterns['call'], line):
-                    call_name = match.group(1)
-                    # Skip common keywords and the current function
-                    if call_name not in ['if', 'for', 'while', 'return', 'print', 'len', 'str', 'int', 'float', 'list', 'dict', 'set', 'tuple', current_scope]:
-                        relationships.append({
-                            'source_scope': current_scope,
-                            'target': call_name,
-                            'type': 'calls',
-                            'file_path': str(file_path),
-                            'line': line_num,
-                        })
+        if TreeSitterSymbolParser is not None:
+            try:
+                ts_parser = TreeSitterSymbolParser(lang, file_path)
+                if ts_parser.is_available():
+                    indexed = ts_parser.parse(content, file_path)
+                    if indexed is not None and indexed.relationships:
+                        relationships = [
+                            {
+                                "source_scope": r.source_symbol,
+                                "target": r.target_symbol,
+                                "type": r.relationship_type.value,
+                                "file_path": str(file_path),
+                                "line": r.source_line,
+                            }
+                            for r in indexed.relationships
+                        ]
+            except Exception:
+                relationships = []
+
+        # Regex fallback for relationships (when tree-sitter is unavailable)
+        if not relationships:
+            current_scope = None
+            for line_num, line in enumerate(lines, 1):
+                for kind in ['function', 'class']:
+                    if kind in patterns:
+                        match = re.search(patterns[kind], line)
+                        if match:
+                            current_scope = match.group(1)
+
+                # Extract imports
+                if 'import' in patterns:
+                    match = re.search(patterns['import'], line)
+                    if match:
+                        import_target = match.group(1) or match.group(2) if match.lastindex >= 2 else match.group(1)
+                        if import_target and current_scope:
+                            relationships.append({
+                                'source_scope': current_scope,
+                                'target': import_target.strip(),
+                                'type': 'imports',
+                                'file_path': str(file_path),
+                                'line': line_num,
+                            })
+
+                # Extract function calls (simplified)
+                if 'call' in patterns and current_scope:
+                    for match in re.finditer(patterns['call'], line):
+                        call_name = match.group(1)
+                        # Skip common keywords and the current function
+                        if call_name not in ['if', 'for', 'while', 'return', 'print', 'len', 'str', 'int', 'float', 'list', 'dict', 'set', 'tuple', current_scope]:
+                            relationships.append({
+                                'source_scope': current_scope,
+                                'target': call_name,
+                                'type': 'calls',
+                                'file_path': str(file_path),
+                                'line': line_num,
+                            })
 
         return symbols, relationships