claude-code-workflow 6.3.18 → 6.3.19

package/codex-lens/src/codexlens/search/ranking.py

@@ -1,7 +1,7 @@
 """Ranking algorithms for hybrid search result fusion.
 
 Implements Reciprocal Rank Fusion (RRF) and score normalization utilities
-for combining results from heterogeneous search backends (exact FTS, fuzzy FTS, vector search).
+for combining results from heterogeneous search backends (SPLADE, exact FTS, fuzzy FTS, vector search).
 """
 
 from __future__ import annotations
@@ -9,11 +9,26 @@ from __future__ import annotations
 import re
 import math
 from enum import Enum
-from typing import Any, Dict, List
+from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 from codexlens.entities import SearchResult, AdditionalLocation
 
 
+# Default RRF weights for SPLADE-based hybrid search
+DEFAULT_WEIGHTS = {
+    "splade": 0.4,  # Replaces exact(0.3) + fuzzy(0.1)
+    "vector": 0.6,
+}
+
+# Legacy weights for FTS fallback mode (when SPLADE unavailable)
+FTS_FALLBACK_WEIGHTS = {
+    "exact": 0.3,
+    "fuzzy": 0.1,
+    "vector": 0.6,
+}
+
+
 class QueryIntent(str, Enum):
     """Query intent for adaptive RRF weights (Python/TypeScript parity)."""
 
@@ -87,15 +102,24 @@ def adjust_weights_by_intent(
     intent: QueryIntent,
     base_weights: Dict[str, float],
 ) -> Dict[str, float]:
-    """Map intent → weights (kept aligned with TypeScript mapping)."""
+    """Adjust RRF weights based on query intent."""
+    # Check if using SPLADE or FTS mode
+    use_splade = "splade" in base_weights
+    
     if intent == QueryIntent.KEYWORD:
-        target = {"exact": 0.5, "fuzzy": 0.1, "vector": 0.4}
+        if use_splade:
+            target = {"splade": 0.6, "vector": 0.4}
+        else:
+            target = {"exact": 0.5, "fuzzy": 0.1, "vector": 0.4}
     elif intent == QueryIntent.SEMANTIC:
-        target = {"exact": 0.2, "fuzzy": 0.1, "vector": 0.7}
+        if use_splade:
+            target = {"splade": 0.3, "vector": 0.7}
+        else:
+            target = {"exact": 0.2, "fuzzy": 0.1, "vector": 0.7}
     else:
         target = dict(base_weights)
-
-    # Preserve only keys that are present in base_weights (active backends).
+    
+    # Filter to active backends
     keys = list(base_weights.keys())
     filtered = {k: float(target.get(k, 0.0)) for k in keys}
     return normalize_weights(filtered)
@@ -109,6 +133,193 @@ def get_rrf_weights(
     return adjust_weights_by_intent(detect_query_intent(query), base_weights)
 
 
+# File extensions to category mapping for fast lookup
+_EXT_TO_CATEGORY: Dict[str, str] = {
+    # Code extensions
+    ".py": "code", ".js": "code", ".jsx": "code", ".ts": "code", ".tsx": "code",
+    ".java": "code", ".go": "code", ".zig": "code", ".m": "code", ".mm": "code",
+    ".c": "code", ".h": "code", ".cc": "code", ".cpp": "code", ".hpp": "code", ".cxx": "code",
+    ".rs": "code",
+    # Doc extensions
+    ".md": "doc", ".mdx": "doc", ".txt": "doc", ".rst": "doc",
+}
+
+
+def get_file_category(path: str) -> Optional[str]:
+    """Get file category ('code' or 'doc') from path extension.
+
+    Args:
+        path: File path string
+
+    Returns:
+        'code', 'doc', or None if unknown
+    """
+    ext = Path(path).suffix.lower()
+    return _EXT_TO_CATEGORY.get(ext)
+
+
+def filter_results_by_category(
+    results: List[SearchResult],
+    intent: QueryIntent,
+    allow_mixed: bool = True,
+) -> List[SearchResult]:
+    """Filter results by category based on query intent.
+
+    Strategy:
+    - KEYWORD (code intent): Only return code files
+    - SEMANTIC (doc intent): Prefer docs, but allow code if allow_mixed=True
+    - MIXED: Return all results
+
+    Args:
+        results: List of SearchResult objects
+        intent: Query intent from detect_query_intent()
+        allow_mixed: If True, SEMANTIC intent includes code files with lower priority
+
+    Returns:
+        Filtered and re-ranked list of SearchResult objects
+    """
+    if not results or intent == QueryIntent.MIXED:
+        return results
+
+    code_results = []
+    doc_results = []
+    unknown_results = []
+
+    for r in results:
+        category = get_file_category(r.path)
+        if category == "code":
+            code_results.append(r)
+        elif category == "doc":
+            doc_results.append(r)
+        else:
+            unknown_results.append(r)
+
+    if intent == QueryIntent.KEYWORD:
+        # Code intent: return only code files + unknown (might be code)
+        filtered = code_results + unknown_results
+    elif intent == QueryIntent.SEMANTIC:
+        if allow_mixed:
+            # Semantic intent with mixed: docs first, then code
+            filtered = doc_results + code_results + unknown_results
+        else:
+            # Semantic intent strict: only docs
+            filtered = doc_results + unknown_results
+    else:
+        filtered = results
+
+    return filtered
+
+
+def simple_weighted_fusion(
+    results_map: Dict[str, List[SearchResult]],
+    weights: Dict[str, float] = None,
+) -> List[SearchResult]:
+    """Combine search results using simple weighted sum of normalized scores.
+
+    This is an alternative to RRF that preserves score magnitude information.
+    Scores are min-max normalized per source before weighted combination.
+
+    Formula: score(d) = Σ weight_source * normalized_score_source(d)
+
+    Args:
+        results_map: Dictionary mapping source name to list of SearchResult objects
+                     Sources: 'exact', 'fuzzy', 'vector', 'splade'
+        weights: Dictionary mapping source name to weight (default: equal weights)
+                 Example: {'exact': 0.3, 'fuzzy': 0.1, 'vector': 0.6}
+
+    Returns:
+        List of SearchResult objects sorted by fused score (descending)
+
+    Examples:
+        >>> fts_results = [SearchResult(path="a.py", score=10.0, excerpt="...")]
+        >>> vector_results = [SearchResult(path="b.py", score=0.85, excerpt="...")]
+        >>> results_map = {'exact': fts_results, 'vector': vector_results}
+        >>> fused = simple_weighted_fusion(results_map)
+    """
+    if not results_map:
+        return []
+
+    # Default equal weights if not provided
+    if weights is None:
+        num_sources = len(results_map)
+        weights = {source: 1.0 / num_sources for source in results_map}
+
+    # Normalize weights to sum to 1.0
+    weight_sum = sum(weights.values())
+    if not math.isclose(weight_sum, 1.0, abs_tol=0.01) and weight_sum > 0:
+        weights = {source: w / weight_sum for source, w in weights.items()}
+
+    # Compute min-max normalization parameters per source
+    source_stats: Dict[str, tuple] = {}
+    for source_name, results in results_map.items():
+        if not results:
+            continue
+        scores = [r.score for r in results]
+        min_s, max_s = min(scores), max(scores)
+        source_stats[source_name] = (min_s, max_s)
+
+    def normalize_score(score: float, source: str) -> float:
+        """Normalize score to [0, 1] range using min-max scaling."""
+        if source not in source_stats:
+            return 0.0
+        min_s, max_s = source_stats[source]
+        if max_s == min_s:
+            return 1.0 if score >= min_s else 0.0
+        return (score - min_s) / (max_s - min_s)
+
+    # Build unified result set with weighted scores
+    path_to_result: Dict[str, SearchResult] = {}
+    path_to_fusion_score: Dict[str, float] = {}
+    path_to_source_scores: Dict[str, Dict[str, float]] = {}
+
+    for source_name, results in results_map.items():
+        weight = weights.get(source_name, 0.0)
+        if weight == 0:
+            continue
+
+        for result in results:
+            path = result.path
+            normalized = normalize_score(result.score, source_name)
+            contribution = weight * normalized
+
+            if path not in path_to_fusion_score:
+                path_to_fusion_score[path] = 0.0
+                path_to_result[path] = result
+                path_to_source_scores[path] = {}
+
+            path_to_fusion_score[path] += contribution
+            path_to_source_scores[path][source_name] = normalized
+
+    # Create final results with fusion scores
+    fused_results = []
+    for path, base_result in path_to_result.items():
+        fusion_score = path_to_fusion_score[path]
+
+        fused_result = SearchResult(
+            path=base_result.path,
+            score=fusion_score,
+            excerpt=base_result.excerpt,
+            content=base_result.content,
+            symbol=base_result.symbol,
+            chunk=base_result.chunk,
+            metadata={
+                **base_result.metadata,
+                "fusion_method": "simple_weighted",
+                "fusion_score": fusion_score,
+                "original_score": base_result.score,
+                "source_scores": path_to_source_scores[path],
+            },
+            start_line=base_result.start_line,
+            end_line=base_result.end_line,
+            symbol_name=base_result.symbol_name,
+            symbol_kind=base_result.symbol_kind,
+        )
+        fused_results.append(fused_result)
+
+    fused_results.sort(key=lambda r: r.score, reverse=True)
+    return fused_results
+
+
 def reciprocal_rank_fusion(
     results_map: Dict[str, List[SearchResult]],
     weights: Dict[str, float] = None,
@@ -118,11 +329,14 @@ def reciprocal_rank_fusion(
 
     RRF formula: score(d) = Σ weight_source / (k + rank_source(d))
 
+    Supports three-way fusion with FTS, Vector, and SPLADE sources.
+
     Args:
         results_map: Dictionary mapping source name to list of SearchResult objects
-                     Sources: 'exact', 'fuzzy', 'vector'
+                     Sources: 'exact', 'fuzzy', 'vector', 'splade'
         weights: Dictionary mapping source name to weight (default: equal weights)
                  Example: {'exact': 0.3, 'fuzzy': 0.1, 'vector': 0.6}
+                 Or: {'splade': 0.4, 'vector': 0.6}
         k: Constant to avoid division by zero and control rank influence (default 60)
 
     Returns:
@@ -133,6 +347,14 @@ def reciprocal_rank_fusion(
         >>> fuzzy_results = [SearchResult(path="b.py", score=8.0, excerpt="...")]
         >>> results_map = {'exact': exact_results, 'fuzzy': fuzzy_results}
         >>> fused = reciprocal_rank_fusion(results_map)
+
+        # Three-way fusion with SPLADE
+        >>> results_map = {
+        ...     'exact': exact_results,
+        ...     'vector': vector_results,
+        ...     'splade': splade_results
+        ... }
+        >>> fused = reciprocal_rank_fusion(results_map, k=60)
     """
     if not results_map:
         return []
@@ -151,6 +373,7 @@ def reciprocal_rank_fusion(
     # Build unified result set with RRF scores
     path_to_result: Dict[str, SearchResult] = {}
     path_to_fusion_score: Dict[str, float] = {}
+    path_to_source_ranks: Dict[str, Dict[str, int]] = {}
 
     for source_name, results in results_map.items():
         weight = weights.get(source_name, 0.0)
@@ -165,8 +388,10 @@ def reciprocal_rank_fusion(
             if path not in path_to_fusion_score:
                 path_to_fusion_score[path] = 0.0
                 path_to_result[path] = result
+                path_to_source_ranks[path] = {}
 
             path_to_fusion_score[path] += rrf_contribution
+            path_to_source_ranks[path][source_name] = rank
 
     # Create final results with fusion scores
     fused_results = []
@@ -183,8 +408,11 @@ def reciprocal_rank_fusion(
             chunk=base_result.chunk,
             metadata={
                 **base_result.metadata,
+                "fusion_method": "rrf",
                 "fusion_score": fusion_score,
                 "original_score": base_result.score,
+                "rrf_k": k,
+                "source_ranks": path_to_source_ranks[path],
             },
             start_line=base_result.start_line,
             end_line=base_result.end_line,
@@ -379,6 +607,117 @@ def rerank_results(
     return reranked_results
 
 
+def cross_encoder_rerank(
+    query: str,
+    results: List[SearchResult],
+    reranker: Any,
+    top_k: int = 50,
+    batch_size: int = 32,
+) -> List[SearchResult]:
+    """Second-stage reranking using a cross-encoder model.
+
+    This function is dependency-agnostic: callers can pass any object that exposes
+    a compatible `score_pairs(pairs, batch_size=...)` method.
+    """
+    if not results:
+        return []
+
+    if reranker is None or top_k <= 0:
+        return results
+
+    rerank_count = min(int(top_k), len(results))
+
+    def text_for_pair(r: SearchResult) -> str:
+        if r.excerpt and r.excerpt.strip():
+            return r.excerpt
+        if r.content and r.content.strip():
+            return r.content
+        if r.chunk and r.chunk.content and r.chunk.content.strip():
+            return r.chunk.content
+        return r.symbol_name or r.path
+
+    pairs = [(query, text_for_pair(r)) for r in results[:rerank_count]]
+
+    try:
+        if hasattr(reranker, "score_pairs"):
+            raw_scores = reranker.score_pairs(pairs, batch_size=int(batch_size))
+        elif hasattr(reranker, "predict"):
+            raw_scores = reranker.predict(pairs, batch_size=int(batch_size))
+        else:
+            return results
+    except Exception:
+        return results
+
+    if not raw_scores or len(raw_scores) != rerank_count:
+        return results
+
+    scores = [float(s) for s in raw_scores]
+    min_s = min(scores)
+    max_s = max(scores)
+
+    def sigmoid(x: float) -> float:
+        # Clamp to keep exp() stable.
+        x = max(-50.0, min(50.0, x))
+        return 1.0 / (1.0 + math.exp(-x))
+
+    if 0.0 <= min_s and max_s <= 1.0:
+        probs = scores
+    else:
+        probs = [sigmoid(s) for s in scores]
+
+    reranked_results: List[SearchResult] = []
+
+    for idx, result in enumerate(results):
+        if idx < rerank_count:
+            prev_score = float(result.score)
+            ce_score = scores[idx]
+            ce_prob = probs[idx]
+            combined_score = 0.5 * prev_score + 0.5 * ce_prob
+
+            reranked_results.append(
+                SearchResult(
+                    path=result.path,
+                    score=combined_score,
+                    excerpt=result.excerpt,
+                    content=result.content,
+                    symbol=result.symbol,
+                    chunk=result.chunk,
+                    metadata={
+                        **result.metadata,
+                        "pre_cross_encoder_score": prev_score,
+                        "cross_encoder_score": ce_score,
+                        "cross_encoder_prob": ce_prob,
+                        "cross_encoder_reranked": True,
+                    },
+                    start_line=result.start_line,
+                    end_line=result.end_line,
+                    symbol_name=result.symbol_name,
+                    symbol_kind=result.symbol_kind,
+                    additional_locations=list(result.additional_locations),
+                )
+            )
+        else:
+            reranked_results.append(
+                SearchResult(
+                    path=result.path,
+                    score=result.score,
+                    excerpt=result.excerpt,
+                    content=result.content,
+                    symbol=result.symbol,
+                    chunk=result.chunk,
+                    metadata={**result.metadata},
+                    start_line=result.start_line,
+                    end_line=result.end_line,
+                    symbol_name=result.symbol_name,
+                    symbol_kind=result.symbol_kind,
+                    additional_locations=list(result.additional_locations),
+                )
+            )
+
+    reranked_results.sort(key=lambda r: r.score, reverse=True)
+    return reranked_results
+
+
 def normalize_bm25_score(score: float) -> float:
     """Normalize BM25 scores from SQLite FTS5 to 0-1 range.
 

package/codex-lens/src/codexlens/semantic/SPLADE_IMPLEMENTATION.md

@@ -0,0 +1,225 @@
+# SPLADE Encoder Implementation
+
+## Overview
+
+Created `splade_encoder.py` - A complete ONNX-optimized SPLADE sparse encoder for code search.
+
+## File Location
+
+`src/codexlens/semantic/splade_encoder.py` (474 lines)
+
+## Key Components
+
+### 1. Dependency Checking
+
+**Function**: `check_splade_available() -> Tuple[bool, Optional[str]]`
+- Validates numpy, onnxruntime, optimum, transformers availability
+- Returns (True, None) if all dependencies present
+- Returns (False, error_message) with install instructions if missing
+
+### 2. Caching System
+
+**Global Cache**: Thread-safe singleton pattern
+- `_splade_cache: Dict[str, SpladeEncoder]` - Global encoder cache
+- `_cache_lock: threading.RLock()` - Thread safety lock
+
+**Factory Function**: `get_splade_encoder(...) -> SpladeEncoder`
+- Cache key includes: model_name, gpu/cpu, max_length, sparsity_threshold
+- Pre-loads model on first access
+- Returns cached instance on subsequent calls
+
+**Cleanup Function**: `clear_splade_cache() -> None`
+- Releases ONNX resources
+- Clears model and tokenizer references
+- Prevents memory leaks
+
+### 3. SpladeEncoder Class
+
+#### Initialization Parameters
+- `model_name: str` - Default: "naver/splade-cocondenser-ensembledistil"
+- `use_gpu: bool` - Enable GPU acceleration (default: True)
+- `max_length: int` - Max sequence length (default: 512)
+- `sparsity_threshold: float` - Min weight threshold (default: 0.01)
+- `providers: Optional[List]` - Explicit ONNX providers (overrides use_gpu)
+
+#### Core Methods
+
+**`_load_model()`**: Lazy loading with GPU support
+- Uses `optimum.onnxruntime.ORTModelForMaskedLM`
+- Falls back to CPU if GPU unavailable
+- Integrates with `gpu_support.get_optimal_providers()`
+- Handles device_id options for DirectML/CUDA
+
+**`_splade_activation(logits, attention_mask)`**: Static method
+- Formula: `log(1 + ReLU(logits)) * attention_mask`
+- Input: (batch, seq_len, vocab_size)
+- Output: (batch, seq_len, vocab_size)
+
+**`_max_pooling(splade_repr)`**: Static method
+- Max pooling over sequence dimension
+- Input: (batch, seq_len, vocab_size)
+- Output: (batch, vocab_size)
+
+**`_to_sparse_dict(dense_vec)`**: Conversion helper
+- Filters by sparsity_threshold
+- Returns: `Dict[int, float]` mapping token_id to weight
+
+**`encode_text(text: str) -> Dict[int, float]`**: Single text encoding
+- Tokenizes input with truncation/padding
+- Forward pass through ONNX model
+- Applies SPLADE activation + max pooling
+- Returns sparse vector
+
+**`encode_batch(texts: List[str], batch_size: int = 32) -> List[Dict[int, float]]`**: Batch encoding
+- Processes in batches for memory efficiency
+- Returns list of sparse vectors
+
+#### Properties
+
+**`vocab_size: int`**: Vocabulary size (~30k for BERT)
+- Cached after first model load
+- Returns tokenizer length
+
+#### Debugging Methods
+
+**`get_token(token_id: int) -> str`**
+- Converts token_id to human-readable string
+- Uses tokenizer.decode()
+
+**`get_top_tokens(sparse_vec: Dict[int, float], top_k: int = 10) -> List[Tuple[str, float]]`**
+- Extracts top-k highest-weight tokens
+- Returns (token_string, weight) pairs
+- Useful for understanding model focus
+
+## Design Patterns Followed
+
+### 1. From `onnx_reranker.py`
+✓ ONNX loading with provider detection
+✓ Lazy model initialization
+✓ Thread-safe loading with RLock
+✓ Signature inspection for backward compatibility
+✓ Fallback for older Optimum versions
+✓ Static helper methods for numerical operations
+
+### 2. From `embedder.py`
+✓ Global cache with thread safety
+✓ Factory function pattern (get_splade_encoder)
+✓ Cache cleanup function (clear_splade_cache)
+✓ GPU provider configuration
+✓ Batch processing support
+
+### 3. From `gpu_support.py`
+✓ `get_optimal_providers(use_gpu, with_device_options=True)`
+✓ Device ID options for DirectML/CUDA
+✓ Provider tuple format: (provider_name, options_dict)
+
+## SPLADE Algorithm
+
+### Activation Formula
+```python
+# Step 1: ReLU activation
+relu_logits = max(0, logits)
+
+# Step 2: Log(1 + x) transformation
+log_relu = log(1 + relu_logits)
+
+# Step 3: Apply attention mask
+splade_repr = log_relu * attention_mask
+
+# Step 4: Max pooling over sequence
+splade_vec = max(splade_repr, axis=sequence_length)
+
+# Step 5: Sparsification by threshold
+sparse_dict = {token_id: weight for token_id, weight in enumerate(splade_vec) if weight > threshold}
+```
+
+### Output Format
+- Sparse dictionary: `{token_id: weight}`
+- Token IDs: 0 to vocab_size-1 (typically ~30,000)
+- Weights: Float values > sparsity_threshold
+- Interpretable: Can decode token_ids to strings
+
+## Integration Points
+
+### With `splade_index.py`
+- `SpladeIndex.add_posting(chunk_id, sparse_vec: Dict[int, float])`
+- `SpladeIndex.search(query_sparse: Dict[int, float])`
+- Encoder produces the sparse vectors consumed by index
+
+### With Indexing Pipeline
+```python
+encoder = get_splade_encoder(use_gpu=True)
+
+# Single document
+sparse_vec = encoder.encode_text("def main():\n    print('hello')")
+index.add_posting(chunk_id=1, sparse_vec=sparse_vec)
+
+# Batch indexing
+texts = ["code chunk 1", "code chunk 2", ...]
+sparse_vecs = encoder.encode_batch(texts, batch_size=64)
+postings = [(chunk_id, vec) for chunk_id, vec in enumerate(sparse_vecs)]
+index.add_postings_batch(postings)
+```
+
+### With Search Pipeline
+```python
+encoder = get_splade_encoder(use_gpu=True)
+query_sparse = encoder.encode_text("authentication function")
+results = index.search(query_sparse, limit=50, min_score=0.5)
+```
+
+## Dependencies
+
+Required packages:
+- `numpy` - Numerical operations
+- `onnxruntime` - ONNX model execution (CPU)
+- `onnxruntime-gpu` - ONNX with GPU support (optional)
+- `optimum[onnxruntime]` - Hugging Face ONNX optimization
+- `transformers` - Tokenizer and model loading
+
+Install command:
+```bash
+# CPU only
+pip install numpy onnxruntime optimum[onnxruntime] transformers
+
+# With GPU support
+pip install numpy onnxruntime-gpu optimum[onnxruntime-gpu] transformers
+```
+
+## Testing Status
+
+✓ Python syntax validation passed
+✓ Module import successful
+✓ Dependency checking works correctly
+✗ Full functional test pending (requires optimum installation)
+
+## Next Steps
+
+1. Install dependencies for functional testing
+2. Create unit tests in `tests/semantic/test_splade_encoder.py`
+3. Benchmark encoding performance (CPU vs GPU)
+4. Integrate with codex-lens indexing pipeline
+5. Add SPLADE option to semantic search configuration
+
+## Performance Considerations
+
+### Memory Usage
+- Model size: ~100MB (ONNX optimized)
+- Sparse vectors: ~100-500 non-zero entries per document
+- Batch size: 32 recommended (adjust based on GPU memory)
+
+### Speed Benchmarks (Expected)
+- CPU encoding: ~10-20 docs/sec
+- GPU encoding (CUDA): ~100-200 docs/sec
+- GPU encoding (DirectML): ~50-100 docs/sec
+
+### Sparsity Analysis
+- Threshold 0.01: ~200-400 tokens per document
+- Threshold 0.05: ~100-200 tokens per document
+- Threshold 0.10: ~50-100 tokens per document
+
+## References
+
+- SPLADE paper: https://arxiv.org/abs/2107.05720
+- SPLADE v2: https://arxiv.org/abs/2109.10086
+- Naver model: https://huggingface.co/naver/splade-cocondenser-ensembledistil