claude-code-workflow 6.2.7 → 6.3.0

package/codex-lens/src/codexlens/semantic/embedder.py

@@ -14,7 +14,8 @@ from typing import Dict, Iterable, List, Optional
 import numpy as np
 
 from . import SEMANTIC_AVAILABLE
-from .gpu_support import get_optimal_providers, is_gpu_available, get_gpu_summary
+from .base import BaseEmbedder
+from .gpu_support import get_optimal_providers, is_gpu_available, get_gpu_summary, get_selected_device_id
 
 logger = logging.getLogger(__name__)
 
@@ -84,7 +85,7 @@ def clear_embedder_cache() -> None:
         gc.collect()
 
 
-class Embedder:
+class Embedder(BaseEmbedder):
     """Generate embeddings for code chunks using fastembed (ONNX-based).
 
     Supported Model Profiles:
@@ -138,25 +139,58 @@ class Embedder:
 
         # Resolve model name from profile or use explicit name
         if model_name:
-            self.model_name = model_name
+            self._model_name = model_name
         elif profile and profile in self.MODELS:
-            self.model_name = self.MODELS[profile]
+            self._model_name = self.MODELS[profile]
         else:
-            self.model_name = self.DEFAULT_MODEL
+            self._model_name = self.DEFAULT_MODEL
 
-        # Configure ONNX execution providers
+        # Configure ONNX execution providers with device_id options for GPU selection
+        # Using with_device_options=True ensures DirectML/CUDA device_id is passed correctly
         if providers is not None:
             self._providers = providers
         else:
-            self._providers = get_optimal_providers(use_gpu=use_gpu)
+            self._providers = get_optimal_providers(use_gpu=use_gpu, with_device_options=True)
 
         self._use_gpu = use_gpu
         self._model = None
 
+    @property
+    def model_name(self) -> str:
+        """Get model name."""
+        return self._model_name
+
     @property
     def embedding_dim(self) -> int:
         """Get embedding dimension for current model."""
-        return self.MODEL_DIMS.get(self.model_name, 768)  # Default to 768 if unknown
+        return self.MODEL_DIMS.get(self._model_name, 768)  # Default to 768 if unknown
+
+    @property
+    def max_tokens(self) -> int:
+        """Get maximum token limit for current model.
+
+        Returns:
+            int: Maximum number of tokens based on model profile.
+                - fast: 512 (lightweight, optimized for speed)
+                - code: 8192 (code-optimized, larger context)
+                - multilingual: 512 (standard multilingual model)
+                - balanced: 512 (general purpose)
+        """
+        # Determine profile from model name
+        profile = None
+        for prof, model in self.MODELS.items():
+            if model == self._model_name:
+                profile = prof
+                break
+
+        # Return token limit based on profile
+        if profile == "code":
+            return 8192
+        elif profile in ("fast", "multilingual", "balanced"):
+            return 512
+        else:
+            # Default for unknown models
+            return 512
 
     @property
     def providers(self) -> List[str]:
@@ -168,7 +202,12 @@ class Embedder:
         """Check if GPU acceleration is enabled for this embedder."""
         gpu_providers = {"CUDAExecutionProvider", "TensorrtExecutionProvider",
                         "DmlExecutionProvider", "ROCMExecutionProvider", "CoreMLExecutionProvider"}
-        return any(p in gpu_providers for p in self._providers)
+        # Handle both string providers and tuple providers (name, options)
+        for p in self._providers:
+            provider_name = p[0] if isinstance(p, tuple) else p
+            if provider_name in gpu_providers:
+                return True
+        return False
 
     def _load_model(self) -> None:
         """Lazy load the embedding model with configured providers."""
@@ -177,7 +216,9 @@ class Embedder:
 
         from fastembed import TextEmbedding
 
-        # fastembed supports 'providers' parameter for ONNX execution providers
+        # providers already include device_id options via get_optimal_providers(with_device_options=True)
+        # DO NOT pass device_ids separately - fastembed ignores it when providers is specified
+        # See: fastembed/text/onnx_embedding.py - device_ids is only used with cuda=True
         try:
             self._model = TextEmbedding(
                 model_name=self.model_name,
@@ -215,7 +256,7 @@ class Embedder:
         embeddings = list(self._model.embed(texts))
         return [emb.tolist() for emb in embeddings]
 
-    def embed_to_numpy(self, texts: str | Iterable[str]) -> np.ndarray:
+    def embed_to_numpy(self, texts: str | Iterable[str], batch_size: Optional[int] = None) -> np.ndarray:
         """Generate embeddings for one or more texts (returns numpy arrays).
 
         This method is more memory-efficient than embed() as it avoids converting
@@ -224,6 +265,8 @@ class Embedder:
 
         Args:
             texts: Single text or iterable of texts to embed.
+            batch_size: Optional batch size for fastembed processing.
+                       Larger values improve GPU utilization but use more memory.
 
         Returns:
             numpy.ndarray of shape (n_texts, embedding_dim) containing embeddings.
@@ -235,8 +278,12 @@ class Embedder:
         else:
             texts = list(texts)
 
-        # Return embeddings as numpy array directly (no .tolist() conversion)
-        embeddings = list(self._model.embed(texts))
+        # Pass batch_size to fastembed for optimal GPU utilization
+        # Default batch_size in fastembed is 256, but larger values can improve throughput
+        if batch_size is not None:
+            embeddings = list(self._model.embed(texts, batch_size=batch_size))
+        else:
+            embeddings = list(self._model.embed(texts))
         return np.array(embeddings)
 
     def embed_single(self, text: str) -> List[float]:

package/codex-lens/src/codexlens/semantic/factory.py

@@ -0,0 +1,98 @@
+"""Factory for creating embedders.
+
+Provides a unified interface for instantiating different embedder backends.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from .base import BaseEmbedder
+
+
+def get_embedder(
+    backend: str = "fastembed",
+    profile: str = "code",
+    model: str = "default",
+    use_gpu: bool = True,
+    endpoints: Optional[List[Dict[str, Any]]] = None,
+    strategy: str = "latency_aware",
+    cooldown: float = 60.0,
+    **kwargs: Any,
+) -> BaseEmbedder:
+    """Factory function to create embedder based on backend.
+
+    Args:
+        backend: Embedder backend to use. Options:
+            - "fastembed": Use fastembed (ONNX-based) embedder (default)
+            - "litellm": Use ccw-litellm embedder
+        profile: Model profile for fastembed backend ("fast", "code", "multilingual", "balanced")
+                Used only when backend="fastembed". Default: "code"
+        model: Model identifier for litellm backend.
+              Used only when backend="litellm". Default: "default"
+        use_gpu: Whether to use GPU acceleration when available (default: True).
+                Used only when backend="fastembed".
+        endpoints: Optional list of endpoint configurations for multi-endpoint load balancing.
+                  Each endpoint is a dict with keys: model, api_key, api_base, weight.
+                  Used only when backend="litellm" and multiple endpoints provided.
+        strategy: Selection strategy for multi-endpoint mode:
+                 "round_robin", "latency_aware", "weighted_random".
+                 Default: "latency_aware"
+        cooldown: Default cooldown seconds for rate-limited endpoints (default: 60.0)
+        **kwargs: Additional backend-specific arguments
+
+    Returns:
+        BaseEmbedder: Configured embedder instance
+
+    Raises:
+        ValueError: If backend is not recognized
+        ImportError: If required backend dependencies are not installed
+
+    Examples:
+        Create fastembed embedder with code profile:
+            >>> embedder = get_embedder(backend="fastembed", profile="code")
+
+        Create fastembed embedder with fast profile and CPU only:
+            >>> embedder = get_embedder(backend="fastembed", profile="fast", use_gpu=False)
+
+        Create litellm embedder:
+            >>> embedder = get_embedder(backend="litellm", model="text-embedding-3-small")
+
+        Create rotational embedder with multiple endpoints:
+            >>> endpoints = [
+            ...     {"model": "openai/text-embedding-3-small", "api_key": "sk-..."},
+            ...     {"model": "azure/my-embedding", "api_base": "https://...", "api_key": "..."},
+            ... ]
+            >>> embedder = get_embedder(backend="litellm", endpoints=endpoints)
+    """
+    if backend == "fastembed":
+        from .embedder import Embedder
+        return Embedder(profile=profile, use_gpu=use_gpu, **kwargs)
+    elif backend == "litellm":
+        # Check if multi-endpoint mode is requested
+        if endpoints and len(endpoints) > 1:
+            from .rotational_embedder import create_rotational_embedder
+            return create_rotational_embedder(
+                endpoints_config=endpoints,
+                strategy=strategy,
+                default_cooldown=cooldown,
+            )
+        elif endpoints and len(endpoints) == 1:
+            # Single endpoint in list - use it directly
+            ep = endpoints[0]
+            ep_kwargs = {**kwargs}
+            if "api_key" in ep:
+                ep_kwargs["api_key"] = ep["api_key"]
+            if "api_base" in ep:
+                ep_kwargs["api_base"] = ep["api_base"]
+            from .litellm_embedder import LiteLLMEmbedderWrapper
+            return LiteLLMEmbedderWrapper(model=ep["model"], **ep_kwargs)
+        else:
+            # No endpoints list - use model parameter
+            from .litellm_embedder import LiteLLMEmbedderWrapper
+            return LiteLLMEmbedderWrapper(model=model, **kwargs)
+    else:
+        raise ValueError(
+            f"Unknown backend: {backend}. "
+            f"Supported backends: 'fastembed', 'litellm'"
+        )

package/codex-lens/src/codexlens/semantic/gpu_support.py

@@ -13,6 +13,15 @@ from typing import List, Optional
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class GPUDevice:
+    """Individual GPU device info."""
+    device_id: int
+    name: str
+    is_discrete: bool  # True for discrete GPU (NVIDIA, AMD), False for integrated (Intel UHD)
+    vendor: str  # "nvidia", "amd", "intel", "unknown"
+
+
 @dataclass
 class GPUInfo:
     """GPU availability and configuration info."""
@@ -22,15 +31,117 @@ class GPUInfo:
     gpu_count: int = 0
     gpu_name: Optional[str] = None
     onnx_providers: List[str] = None
+    devices: List[GPUDevice] = None  # List of detected GPU devices
+    preferred_device_id: Optional[int] = None  # Preferred GPU for embedding
 
     def __post_init__(self):
         if self.onnx_providers is None:
             self.onnx_providers = ["CPUExecutionProvider"]
+        if self.devices is None:
+            self.devices = []
 
 
 _gpu_info_cache: Optional[GPUInfo] = None
 
 
+def _enumerate_gpus() -> List[GPUDevice]:
+    """Enumerate available GPU devices using WMI on Windows.
+    
+    Returns:
+        List of GPUDevice with device info, ordered by device_id.
+    """
+    devices = []
+    
+    try:
+        import subprocess
+        import sys
+        
+        if sys.platform == "win32":
+            # Use PowerShell to query GPU information via WMI
+            cmd = [
+                "powershell", "-NoProfile", "-Command",
+                "Get-WmiObject Win32_VideoController | Select-Object DeviceID, Name, AdapterCompatibility | ConvertTo-Json"
+            ]
+            result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
+            
+            if result.returncode == 0 and result.stdout.strip():
+                import json
+                gpu_data = json.loads(result.stdout)
+                
+                # Handle single GPU case (returns dict instead of list)
+                if isinstance(gpu_data, dict):
+                    gpu_data = [gpu_data]
+                
+                for idx, gpu in enumerate(gpu_data):
+                    name = gpu.get("Name", "Unknown GPU")
+                    compat = gpu.get("AdapterCompatibility", "").lower()
+                    
+                    # Determine vendor
+                    name_lower = name.lower()
+                    if "nvidia" in name_lower or "nvidia" in compat:
+                        vendor = "nvidia"
+                        is_discrete = True
+                    elif "amd" in name_lower or "radeon" in name_lower or "amd" in compat:
+                        vendor = "amd"
+                        is_discrete = True
+                    elif "intel" in name_lower or "intel" in compat:
+                        vendor = "intel"
+                        # Intel UHD/Iris are integrated, Intel Arc is discrete
+                        is_discrete = "arc" in name_lower
+                    else:
+                        vendor = "unknown"
+                        is_discrete = False
+                    
+                    devices.append(GPUDevice(
+                        device_id=idx,
+                        name=name,
+                        is_discrete=is_discrete,
+                        vendor=vendor
+                    ))
+                    logger.debug(f"Detected GPU {idx}: {name} (vendor={vendor}, discrete={is_discrete})")
+                    
+    except Exception as e:
+        logger.debug(f"GPU enumeration failed: {e}")
+    
+    return devices
+
+
+def _get_preferred_device_id(devices: List[GPUDevice]) -> Optional[int]:
+    """Determine the preferred GPU device_id for embedding.
+    
+    Preference order:
+    1. NVIDIA discrete GPU (best DirectML/CUDA support)
+    2. AMD discrete GPU
+    3. Intel Arc (discrete)
+    4. Intel integrated (fallback)
+    
+    Returns:
+        device_id of preferred GPU, or None to use default.
+    """
+    if not devices:
+        return None
+    
+    # Priority: NVIDIA > AMD > Intel Arc > Intel integrated
+    priority_order = [
+        ("nvidia", True),   # NVIDIA discrete
+        ("amd", True),      # AMD discrete
+        ("intel", True),    # Intel Arc (discrete)
+        ("intel", False),   # Intel integrated (fallback)
+    ]
+    
+    for target_vendor, target_discrete in priority_order:
+        for device in devices:
+            if device.vendor == target_vendor and device.is_discrete == target_discrete:
+                logger.info(f"Preferred GPU: {device.name} (device_id={device.device_id})")
+                return device.device_id
+    
+    # If no match, use first device
+    if devices:
+        return devices[0].device_id
+    
+    return None
+
+
 def detect_gpu(force_refresh: bool = False) -> GPUInfo:
     """Detect available GPU resources for embedding acceleration.
 
@@ -47,6 +158,18 @@ def detect_gpu(force_refresh: bool = False) -> GPUInfo:
 
     info = GPUInfo()
 
+    # Enumerate GPU devices first
+    info.devices = _enumerate_gpus()
+    info.gpu_count = len(info.devices)
+    if info.devices:
+        # Set preferred device (discrete GPU preferred over integrated)
+        info.preferred_device_id = _get_preferred_device_id(info.devices)
+        # Set gpu_name to preferred device name
+        for dev in info.devices:
+            if dev.device_id == info.preferred_device_id:
+                info.gpu_name = dev.name
+                break
+
     # Check PyTorch CUDA availability (most reliable detection)
     try:
         import torch
@@ -143,21 +266,48 @@ def detect_gpu(force_refresh: bool = False) -> GPUInfo:
     return info
 
 
-def get_optimal_providers(use_gpu: bool = True) -> List[str]:
+def get_optimal_providers(use_gpu: bool = True, with_device_options: bool = False) -> list:
     """Get optimal ONNX execution providers based on availability.
 
     Args:
         use_gpu: If True, include GPU providers when available.
                  If False, force CPU-only execution.
+        with_device_options: If True, return providers as tuples with device_id options
+                            for proper GPU device selection (required for DirectML).
 
     Returns:
-        List of provider names in priority order.
+        List of provider names or tuples (provider_name, options_dict) in priority order.
     """
     if not use_gpu:
         return ["CPUExecutionProvider"]
 
     gpu_info = detect_gpu()
-    return gpu_info.onnx_providers
+    
+    if not with_device_options:
+        return gpu_info.onnx_providers
+    
+    # Build providers with device_id options for GPU providers
+    device_id = get_selected_device_id()
+    providers = []
+    
+    for provider in gpu_info.onnx_providers:
+        if provider == "DmlExecutionProvider" and device_id is not None:
+            # DirectML requires device_id in provider_options tuple
+            providers.append(("DmlExecutionProvider", {"device_id": device_id}))
+            logger.debug(f"DmlExecutionProvider configured with device_id={device_id}")
+        elif provider == "CUDAExecutionProvider" and device_id is not None:
+            # CUDA also supports device_id in provider_options
+            providers.append(("CUDAExecutionProvider", {"device_id": device_id}))
+            logger.debug(f"CUDAExecutionProvider configured with device_id={device_id}")
+        elif provider == "ROCMExecutionProvider" and device_id is not None:
+            # ROCm supports device_id
+            providers.append(("ROCMExecutionProvider", {"device_id": device_id}))
+            logger.debug(f"ROCMExecutionProvider configured with device_id={device_id}")
+        else:
+            # CPU and other providers don't need device_id
+            providers.append(provider)
+    
+    return providers
 
 
 def is_gpu_available() -> bool:
@@ -190,3 +340,75 @@ def clear_gpu_cache() -> None:
     """Clear cached GPU detection info."""
     global _gpu_info_cache
     _gpu_info_cache = None
+
+
+# User-selected device ID (overrides auto-detection)
+_selected_device_id: Optional[int] = None
+
+
+def get_gpu_devices() -> List[dict]:
+    """Get list of available GPU devices for frontend selection.
+    
+    Returns:
+        List of dicts with device info for each GPU.
+    """
+    info = detect_gpu()
+    devices = []
+    
+    for dev in info.devices:
+        devices.append({
+            "device_id": dev.device_id,
+            "name": dev.name,
+            "vendor": dev.vendor,
+            "is_discrete": dev.is_discrete,
+            "is_preferred": dev.device_id == info.preferred_device_id,
+            "is_selected": dev.device_id == get_selected_device_id(),
+        })
+    
+    return devices
+
+
+def get_selected_device_id() -> Optional[int]:
+    """Get the user-selected GPU device_id.
+    
+    Returns:
+        User-selected device_id, or auto-detected preferred device_id if not set.
+    """
+    global _selected_device_id
+    
+    if _selected_device_id is not None:
+        return _selected_device_id
+    
+    # Fall back to auto-detected preferred device
+    info = detect_gpu()
+    return info.preferred_device_id
+
+
+def set_selected_device_id(device_id: Optional[int]) -> bool:
+    """Set the GPU device_id to use for embeddings.
+    
+    Args:
+        device_id: GPU device_id to use, or None to use auto-detection.
+        
+    Returns:
+        True if device_id is valid, False otherwise.
+    """
+    global _selected_device_id
+    
+    if device_id is None:
+        _selected_device_id = None
+        logger.info("GPU selection reset to auto-detection")
+        return True
+    
+    # Validate device_id exists
+    info = detect_gpu()
+    valid_ids = [dev.device_id for dev in info.devices]
+    
+    if device_id in valid_ids:
+        _selected_device_id = device_id
+        device_name = next((dev.name for dev in info.devices if dev.device_id == device_id), "Unknown")
+        logger.info(f"GPU selection set to device {device_id}: {device_name}")
+        return True
+    else:
+        logger.warning(f"Invalid device_id {device_id}. Valid IDs: {valid_ids}")
+        return False

package/codex-lens/src/codexlens/semantic/litellm_embedder.py

@@ -0,0 +1,144 @@
+"""LiteLLM embedder wrapper for CodexLens.
+
+Provides integration with ccw-litellm's LiteLLMEmbedder for embedding generation.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import numpy as np
+
+from .base import BaseEmbedder
+
+
+class LiteLLMEmbedderWrapper(BaseEmbedder):
+    """Wrapper for ccw-litellm LiteLLMEmbedder.
+
+    This wrapper adapts the ccw-litellm LiteLLMEmbedder to the CodexLens
+    BaseEmbedder interface, enabling seamless integration with CodexLens
+    semantic search functionality.
+
+    Args:
+        model: Model identifier for LiteLLM (default: "default")
+        **kwargs: Additional arguments passed to LiteLLMEmbedder
+
+    Raises:
+        ImportError: If ccw-litellm package is not installed
+    """
+
+    def __init__(self, model: str = "default", **kwargs) -> None:
+        """Initialize LiteLLM embedder wrapper.
+
+        Args:
+            model: Model identifier for LiteLLM (default: "default")
+            **kwargs: Additional arguments passed to LiteLLMEmbedder
+
+        Raises:
+            ImportError: If ccw-litellm package is not installed
+        """
+        try:
+            from ccw_litellm import LiteLLMEmbedder
+            self._embedder = LiteLLMEmbedder(model=model, **kwargs)
+        except ImportError as e:
+            raise ImportError(
+                "ccw-litellm not installed. Install with: pip install ccw-litellm"
+            ) from e
+
+    @property
+    def embedding_dim(self) -> int:
+        """Return embedding dimensions from LiteLLMEmbedder.
+
+        Returns:
+            int: Dimension of the embedding vectors.
+        """
+        return self._embedder.dimensions
+
+    @property
+    def model_name(self) -> str:
+        """Return model name from LiteLLMEmbedder.
+
+        Returns:
+            str: Name or identifier of the underlying model.
+        """
+        return self._embedder.model_name
+
+    @property
+    def max_tokens(self) -> int:
+        """Return maximum token limit for the embedding model.
+
+        Returns:
+            int: Maximum number of tokens that can be embedded at once.
+                Inferred from model config or model name patterns.
+        """
+        # Try to get from LiteLLM config first
+        if hasattr(self._embedder, 'max_input_tokens') and self._embedder.max_input_tokens:
+            return self._embedder.max_input_tokens
+
+        # Infer from model name
+        model_name_lower = self.model_name.lower()
+
+        # Large models (8B or "large" in name)
+        if '8b' in model_name_lower or 'large' in model_name_lower:
+            return 32768
+
+        # OpenAI text-embedding-3-* models
+        if 'text-embedding-3' in model_name_lower:
+            return 8191
+
+        # Default fallback
+        return 8192
+
+    def _sanitize_text(self, text: str) -> str:
+        """Sanitize text to work around ModelScope API routing bug.
+
+        ModelScope incorrectly routes text starting with lowercase 'import'
+        to an Ollama endpoint, causing failures. This adds a leading space
+        to work around the issue without affecting embedding quality.
+
+        Args:
+            text: Text to sanitize.
+
+        Returns:
+            Sanitized text safe for embedding API.
+        """
+        if text.startswith('import'):
+            return ' ' + text
+        return text
+
+    def embed_to_numpy(self, texts: str | Iterable[str], **kwargs) -> np.ndarray:
+        """Embed texts to numpy array using LiteLLMEmbedder.
+
+        Args:
+            texts: Single text or iterable of texts to embed.
+            **kwargs: Additional arguments (ignored for LiteLLM backend).
+                      Accepts batch_size for API compatibility with fastembed.
+
+        Returns:
+            numpy.ndarray: Array of shape (n_texts, embedding_dim) containing embeddings.
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        else:
+            texts = list(texts)
+
+        # Sanitize texts to avoid ModelScope routing bug
+        texts = [self._sanitize_text(t) for t in texts]
+
+        # LiteLLM handles batching internally, ignore batch_size parameter
+        return self._embedder.embed(texts)
+
+    def embed_single(self, text: str) -> list[float]:
+        """Generate embedding for a single text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            list[float]: Embedding vector as a list of floats.
+        """
+        # Sanitize text before embedding
+        sanitized = self._sanitize_text(text)
+        embedding = self._embedder.embed([sanitized])
+        return embedding[0].tolist()
+