claude-code-workflow 6.3.18 → 6.3.20

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,167 @@ 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,
+    chunk_type_weights: Optional[Dict[str, float]] = None,
+    test_file_penalty: float = 0.0,
+) -> 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.
+
+    Args:
+        query: Search query string
+        results: List of search results to rerank
+        reranker: Cross-encoder model with score_pairs or predict method
+        top_k: Number of top results to rerank
+        batch_size: Batch size for reranking
+        chunk_type_weights: Optional weights for different chunk types.
+            Example: {"code": 1.0, "docstring": 0.7} - reduce docstring influence
+        test_file_penalty: Penalty applied to test files (0.0-1.0).
+            Example: 0.2 means test files get 20% score reduction
+    """
+    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] = []
+
+    # Helper to detect test files
+    def is_test_file(path: str) -> bool:
+        if not path:
+            return False
+        basename = path.split("/")[-1].split("\\")[-1]
+        return (
+            basename.startswith("test_") or
+            basename.endswith("_test.py") or
+            basename.endswith(".test.ts") or
+            basename.endswith(".test.js") or
+            basename.endswith(".spec.ts") or
+            basename.endswith(".spec.js") or
+            "/tests/" in path or
+            "\\tests\\" in path or
+            "/test/" in path or
+            "\\test\\" in path
+        )
+
+    for idx, result in enumerate(results):
+        if idx < rerank_count:
+            prev_score = float(result.score)
+            ce_score = scores[idx]
+            ce_prob = probs[idx]
+
+            # Base combined score
+            combined_score = 0.5 * prev_score + 0.5 * ce_prob
+
+            # Apply chunk_type weight adjustment
+            if chunk_type_weights:
+                chunk_type = None
+                if result.chunk and hasattr(result.chunk, "metadata"):
+                    chunk_type = result.chunk.metadata.get("chunk_type")
+                elif result.metadata:
+                    chunk_type = result.metadata.get("chunk_type")
+
+                if chunk_type and chunk_type in chunk_type_weights:
+                    weight = chunk_type_weights[chunk_type]
+                    # Apply weight to CE contribution only
+                    combined_score = 0.5 * prev_score + 0.5 * ce_prob * weight
+
+            # Apply test file penalty
+            if test_file_penalty > 0 and is_test_file(result.path):
+                combined_score = combined_score * (1.0 - test_file_penalty)
+
+            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.