claude-code-workflow 6.3.18 → 6.3.19

package/codex-lens/src/codexlens/semantic/reranker/fastembed_reranker.py

@@ -0,0 +1,257 @@
+"""FastEmbed-based reranker backend.
+
+This reranker uses fastembed's TextCrossEncoder for cross-encoder reranking.
+FastEmbed is ONNX-based internally but provides a cleaner, unified API.
+
+Install:
+    pip install fastembed>=0.4.0
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Sequence
+
+from .base import BaseReranker
+
+logger = logging.getLogger(__name__)
+
+
+def check_fastembed_reranker_available() -> tuple[bool, str | None]:
+    """Check whether fastembed reranker dependencies are available."""
+    try:
+        import fastembed  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        return (
+            False,
+            f"fastembed not available: {exc}. Install with: pip install fastembed>=0.4.0",
+        )
+
+    try:
+        from fastembed.rerank.cross_encoder import TextCrossEncoder  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        return (
+            False,
+            f"fastembed TextCrossEncoder not available: {exc}. "
+            "Upgrade with: pip install fastembed>=0.4.0",
+        )
+
+    return True, None
+
+
+class FastEmbedReranker(BaseReranker):
+    """Cross-encoder reranker using fastembed's TextCrossEncoder with lazy loading."""
+
+    DEFAULT_MODEL = "Xenova/ms-marco-MiniLM-L-6-v2"
+
+    # Alternative models supported by fastembed:
+    # - "BAAI/bge-reranker-base"
+    # - "BAAI/bge-reranker-large"
+    # - "cross-encoder/ms-marco-MiniLM-L-6-v2"
+
+    def __init__(
+        self,
+        model_name: str | None = None,
+        *,
+        use_gpu: bool = True,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+    ) -> None:
+        """Initialize FastEmbed reranker.
+
+        Args:
+            model_name: Model identifier. Defaults to Xenova/ms-marco-MiniLM-L-6-v2.
+            use_gpu: Whether to use GPU acceleration when available.
+            cache_dir: Optional directory for caching downloaded models.
+            threads: Optional number of threads for ONNX Runtime.
+        """
+        self.model_name = (model_name or self.DEFAULT_MODEL).strip()
+        if not self.model_name:
+            raise ValueError("model_name cannot be blank")
+
+        self.use_gpu = bool(use_gpu)
+        self.cache_dir = cache_dir
+        self.threads = threads
+
+        self._encoder: Any | None = None
+        self._lock = threading.RLock()
+
+    def _load_model(self) -> None:
+        """Lazy-load the TextCrossEncoder model."""
+        if self._encoder is not None:
+            return
+
+        ok, err = check_fastembed_reranker_available()
+        if not ok:
+            raise ImportError(err)
+
+        with self._lock:
+            if self._encoder is not None:
+                return
+
+            from fastembed.rerank.cross_encoder import TextCrossEncoder
+
+            # Determine providers based on GPU preference
+            providers: list[str] | None = None
+            if self.use_gpu:
+                try:
+                    from ..gpu_support import get_optimal_providers
+
+                    providers = get_optimal_providers(use_gpu=True, with_device_options=False)
+                except Exception:
+                    # Fallback: let fastembed decide
+                    providers = None
+
+            # Build initialization kwargs
+            init_kwargs: dict[str, Any] = {}
+            if self.cache_dir:
+                init_kwargs["cache_dir"] = self.cache_dir
+            if self.threads is not None:
+                init_kwargs["threads"] = self.threads
+            if providers:
+                init_kwargs["providers"] = providers
+
+            logger.debug(
+                "Loading FastEmbed reranker model: %s (use_gpu=%s)",
+                self.model_name,
+                self.use_gpu,
+            )
+
+            self._encoder = TextCrossEncoder(
+                model_name=self.model_name,
+                **init_kwargs,
+            )
+
+            logger.debug("FastEmbed reranker model loaded successfully")
+
+    @staticmethod
+    def _sigmoid(x: float) -> float:
+        """Numerically stable sigmoid function."""
+        if x < -709:
+            return 0.0
+        if x > 709:
+            return 1.0
+        import math
+        return 1.0 / (1.0 + math.exp(-x))
+
+    def score_pairs(
+        self,
+        pairs: Sequence[tuple[str, str]],
+        *,
+        batch_size: int = 32,
+    ) -> list[float]:
+        """Score (query, doc) pairs.
+
+        Args:
+            pairs: Sequence of (query, doc) string pairs to score.
+            batch_size: Batch size for scoring.
+
+        Returns:
+            List of scores (one per pair), normalized to [0, 1] range.
+        """
+        if not pairs:
+            return []
+
+        self._load_model()
+
+        if self._encoder is None:  # pragma: no cover - defensive
+            return []
+
+        # FastEmbed's TextCrossEncoder.rerank() expects a query and list of documents.
+        # For batch scoring of multiple query-doc pairs, we need to process them.
+        # Group by query for efficiency when same query appears multiple times.
+        query_to_docs: dict[str, list[tuple[int, str]]] = {}
+        for idx, (query, doc) in enumerate(pairs):
+            if query not in query_to_docs:
+                query_to_docs[query] = []
+            query_to_docs[query].append((idx, doc))
+
+        # Score each query group
+        scores: list[float] = [0.0] * len(pairs)
+
+        for query, indexed_docs in query_to_docs.items():
+            docs = [doc for _, doc in indexed_docs]
+            indices = [idx for idx, _ in indexed_docs]
+
+            try:
+                # TextCrossEncoder.rerank returns raw float scores in same order as input
+                raw_scores = list(
+                    self._encoder.rerank(
+                        query=query,
+                        documents=docs,
+                        batch_size=batch_size,
+                    )
+                )
+
+                # Map scores back to original positions and normalize with sigmoid
+                for i, raw_score in enumerate(raw_scores):
+                    if i < len(indices):
+                        original_idx = indices[i]
+                        # Normalize score to [0, 1] using stable sigmoid
+                        scores[original_idx] = self._sigmoid(float(raw_score))
+
+            except Exception as e:
+                logger.warning("FastEmbed rerank failed for query: %s", str(e)[:100])
+                # Leave scores as 0.0 for failed queries
+
+        return scores
+
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[str],
+        *,
+        top_k: int | None = None,
+        batch_size: int = 32,
+    ) -> list[tuple[float, str, int]]:
+        """Rerank documents for a single query.
+
+        This is a convenience method that provides results in ranked order.
+
+        Args:
+            query: The query string.
+            documents: List of documents to rerank.
+            top_k: Return only top K results. None returns all.
+            batch_size: Batch size for scoring.
+
+        Returns:
+            List of (score, document, original_index) tuples, sorted by score descending.
+        """
+        if not documents:
+            return []
+
+        self._load_model()
+
+        if self._encoder is None:  # pragma: no cover - defensive
+            return []
+
+        try:
+            # TextCrossEncoder.rerank returns raw float scores in same order as input
+            raw_scores = list(
+                self._encoder.rerank(
+                    query=query,
+                    documents=list(documents),
+                    batch_size=batch_size,
+                )
+            )
+
+            # Convert to our format: (normalized_score, document, original_index)
+            ranked = []
+            for idx, raw_score in enumerate(raw_scores):
+                if idx < len(documents):
+                    # Normalize score to [0, 1] using stable sigmoid
+                    normalized = self._sigmoid(float(raw_score))
+                    ranked.append((normalized, documents[idx], idx))
+
+            # Sort by score descending
+            ranked.sort(key=lambda x: x[0], reverse=True)
+
+            if top_k is not None and top_k > 0:
+                ranked = ranked[:top_k]
+
+            return ranked
+
+        except Exception as e:
+            logger.warning("FastEmbed rerank failed: %s", str(e)[:100])
+            return []

package/codex-lens/src/codexlens/semantic/reranker/legacy.py

@@ -0,0 +1,91 @@
+"""Legacy sentence-transformers cross-encoder reranker.
+
+Install with: pip install codexlens[reranker-legacy]
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import List, Sequence, Tuple
+
+from .base import BaseReranker
+
+logger = logging.getLogger(__name__)
+
+try:
+    from sentence_transformers import CrossEncoder as _CrossEncoder
+
+    CROSS_ENCODER_AVAILABLE = True
+    _import_error: str | None = None
+except ImportError as exc:  # pragma: no cover - optional dependency
+    _CrossEncoder = None  # type: ignore[assignment]
+    CROSS_ENCODER_AVAILABLE = False
+    _import_error = str(exc)
+
+
+def check_cross_encoder_available() -> tuple[bool, str | None]:
+    if CROSS_ENCODER_AVAILABLE:
+        return True, None
+    return (
+        False,
+        _import_error
+        or "sentence-transformers not available. Install with: pip install codexlens[reranker-legacy]",
+    )
+
+
+class CrossEncoderReranker(BaseReranker):
+    """Cross-encoder reranker with lazy model loading."""
+
+    def __init__(self, model_name: str, *, device: str | None = None) -> None:
+        self.model_name = (model_name or "").strip()
+        if not self.model_name:
+            raise ValueError("model_name cannot be blank")
+
+        self.device = (device or "").strip() or None
+        self._model = None
+        self._lock = threading.RLock()
+
+    def _load_model(self) -> None:
+        if self._model is not None:
+            return
+
+        ok, err = check_cross_encoder_available()
+        if not ok:
+            raise ImportError(err)
+
+        with self._lock:
+            if self._model is not None:
+                return
+
+            try:
+                if self.device:
+                    self._model = _CrossEncoder(self.model_name, device=self.device)  # type: ignore[misc]
+                else:
+                    self._model = _CrossEncoder(self.model_name)  # type: ignore[misc]
+            except Exception as exc:
+                logger.debug("Failed to load cross-encoder model %s: %s", self.model_name, exc)
+                raise
+
+    def score_pairs(
+        self,
+        pairs: Sequence[Tuple[str, str]],
+        *,
+        batch_size: int = 32,
+    ) -> List[float]:
+        """Score (query, doc) pairs using the cross-encoder.
+
+        Returns:
+            List of scores (one per pair) in the model's native scale (usually logits).
+        """
+        if not pairs:
+            return []
+
+        self._load_model()
+
+        if self._model is None:  # pragma: no cover - defensive
+            return []
+
+        bs = int(batch_size) if batch_size and int(batch_size) > 0 else 32
+        scores = self._model.predict(list(pairs), batch_size=bs)  # type: ignore[union-attr]
+        return [float(s) for s in scores]

package/codex-lens/src/codexlens/semantic/reranker/litellm_reranker.py

@@ -0,0 +1,214 @@
+"""Experimental LiteLLM reranker backend.
+
+This module provides :class:`LiteLLMReranker`, which uses an LLM to score the
+relevance of a single (query, document) pair per request.
+
+Notes:
+    - This backend is experimental and may be slow/expensive compared to local
+      rerankers.
+    - It relies on `ccw-litellm` for a unified LLM API across providers.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+import threading
+import time
+from typing import Any, Sequence
+
+from .base import BaseReranker
+
+logger = logging.getLogger(__name__)
+
+_NUMBER_RE = re.compile(r"[-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?")
+
+
+def _coerce_score_to_unit_interval(score: float) -> float:
+    """Coerce a numeric score into [0, 1].
+
+    The prompt asks for a float in [0, 1], but some models may respond with 0-10
+    or 0-100 scales. This function attempts a conservative normalization.
+    """
+    if 0.0 <= score <= 1.0:
+        return score
+    if 0.0 <= score <= 10.0:
+        return score / 10.0
+    if 0.0 <= score <= 100.0:
+        return score / 100.0
+    return max(0.0, min(1.0, score))
+
+
+def _extract_score(text: str) -> float | None:
+    """Extract a numeric relevance score from an LLM response."""
+    content = (text or "").strip()
+    if not content:
+        return None
+
+    # Prefer JSON if present.
+    if "{" in content and "}" in content:
+        try:
+            start = content.index("{")
+            end = content.rindex("}") + 1
+            payload = json.loads(content[start:end])
+            if isinstance(payload, dict) and "score" in payload:
+                return float(payload["score"])
+        except Exception:
+            pass
+
+    match = _NUMBER_RE.search(content)
+    if not match:
+        return None
+    try:
+        return float(match.group(0))
+    except ValueError:
+        return None
+
+
+class LiteLLMReranker(BaseReranker):
+    """Experimental reranker that uses a LiteLLM-compatible model.
+
+    This reranker scores each (query, doc) pair in isolation (single-pair mode)
+    to improve prompt reliability across providers.
+    """
+
+    _SYSTEM_PROMPT = (
+        "You are a relevance scoring assistant.\n"
+        "Given a search query and a document snippet, output a single numeric "
+        "relevance score between 0 and 1.\n\n"
+        "Scoring guidance:\n"
+        "- 1.0: The document directly answers the query.\n"
+        "- 0.5: The document is partially relevant.\n"
+        "- 0.0: The document is unrelated.\n\n"
+        "Output requirements:\n"
+        "- Output ONLY the number (e.g., 0.73).\n"
+        "- Do not include any other text."
+    )
+
+    def __init__(
+        self,
+        model: str = "default",
+        *,
+        requests_per_minute: float | None = None,
+        min_interval_seconds: float | None = None,
+        default_score: float = 0.0,
+        max_doc_chars: int = 8000,
+        **litellm_kwargs: Any,
+    ) -> None:
+        """Initialize the reranker.
+
+        Args:
+            model: Model name from ccw-litellm configuration (default: "default").
+            requests_per_minute: Optional rate limit in requests per minute.
+            min_interval_seconds: Optional minimum interval between requests. If set,
+                it takes precedence over requests_per_minute.
+            default_score: Score to use when an API call fails or parsing fails.
+            max_doc_chars: Maximum number of document characters to include in the prompt.
+            **litellm_kwargs: Passed through to `ccw_litellm.LiteLLMClient`.
+
+        Raises:
+            ImportError: If ccw-litellm is not installed.
+            ValueError: If model is blank.
+        """
+        self.model_name = (model or "").strip()
+        if not self.model_name:
+            raise ValueError("model cannot be blank")
+
+        self.default_score = float(default_score)
+
+        self.max_doc_chars = int(max_doc_chars) if int(max_doc_chars) > 0 else 0
+
+        if min_interval_seconds is not None:
+            self._min_interval_seconds = max(0.0, float(min_interval_seconds))
+        elif requests_per_minute is not None and float(requests_per_minute) > 0:
+            self._min_interval_seconds = 60.0 / float(requests_per_minute)
+        else:
+            self._min_interval_seconds = 0.0
+
+        # Prefer deterministic output by default; allow overrides via kwargs.
+        litellm_kwargs = dict(litellm_kwargs)
+        litellm_kwargs.setdefault("temperature", 0.0)
+        litellm_kwargs.setdefault("max_tokens", 16)
+
+        try:
+            from ccw_litellm import ChatMessage, LiteLLMClient
+        except ImportError as exc:  # pragma: no cover - optional dependency
+            raise ImportError(
+                "ccw-litellm not installed. Install with: pip install ccw-litellm"
+            ) from exc
+
+        self._ChatMessage = ChatMessage
+        self._client = LiteLLMClient(model=self.model_name, **litellm_kwargs)
+
+        self._lock = threading.RLock()
+        self._last_request_at = 0.0
+
+    def _sanitize_text(self, text: str) -> str:
+        # Keep consistent with LiteLLMEmbedderWrapper workaround.
+        if text.startswith("import"):
+            return " " + text
+        return text
+
+    def _rate_limit(self) -> None:
+        if self._min_interval_seconds <= 0:
+            return
+        with self._lock:
+            now = time.monotonic()
+            elapsed = now - self._last_request_at
+            if elapsed < self._min_interval_seconds:
+                time.sleep(self._min_interval_seconds - elapsed)
+            self._last_request_at = time.monotonic()
+
+    def _build_user_prompt(self, query: str, doc: str) -> str:
+        sanitized_query = self._sanitize_text(query or "")
+        sanitized_doc = self._sanitize_text(doc or "")
+        if self.max_doc_chars and len(sanitized_doc) > self.max_doc_chars:
+            sanitized_doc = sanitized_doc[: self.max_doc_chars]
+
+        return (
+            "Query:\n"
+            f"{sanitized_query}\n\n"
+            "Document:\n"
+            f"{sanitized_doc}\n\n"
+            "Return the relevance score (0 to 1) as a single number:"
+        )
+
+    def _score_single_pair(self, query: str, doc: str) -> float:
+        messages = [
+            self._ChatMessage(role="system", content=self._SYSTEM_PROMPT),
+            self._ChatMessage(role="user", content=self._build_user_prompt(query, doc)),
+        ]
+
+        try:
+            self._rate_limit()
+            response = self._client.chat(messages)
+        except Exception as exc:
+            logger.debug("LiteLLM reranker request failed: %s", exc)
+            return self.default_score
+
+        raw = getattr(response, "content", "") or ""
+        score = _extract_score(raw)
+        if score is None:
+            logger.debug("Failed to parse LiteLLM reranker score from response: %r", raw)
+            return self.default_score
+        return _coerce_score_to_unit_interval(float(score))
+
+    def score_pairs(
+        self,
+        pairs: Sequence[tuple[str, str]],
+        *,
+        batch_size: int = 32,
+    ) -> list[float]:
+        """Score (query, doc) pairs with per-pair LLM calls."""
+        if not pairs:
+            return []
+
+        bs = int(batch_size) if batch_size and int(batch_size) > 0 else 32
+
+        scores: list[float] = []
+        for i in range(0, len(pairs), bs):
+            batch = pairs[i : i + bs]
+            for query, doc in batch:
+                scores.append(self._score_single_pair(query, doc))
+        return scores