Merge pull request #180 from xmican10/LEADS-241-JudgeLLM-token-counter-for-Deepeval

[LEADS-241] judge llm token counter for deepeval

Author: asamal4

src/lightspeed_evaluation/core/llm/__init__.py

@@ -2,11 +2,15 @@
 
 from typing import TYPE_CHECKING
 
+# Apply litellm patching globally before any litellm usage in this package
+import lightspeed_evaluation.core.llm.litellm_patch  # noqa: F401
+
 from lightspeed_evaluation.core.system.lazy_import import create_lazy_getattr
 
 if TYPE_CHECKING:
     # ruff: noqa: F401
-    from lightspeed_evaluation.core.llm.custom import BaseCustomLLM, TokenTracker
+    from lightspeed_evaluation.core.llm.custom import BaseCustomLLM
+    from lightspeed_evaluation.core.llm.token_tracker import TokenTracker
     from lightspeed_evaluation.core.llm.deepeval import DeepEvalLLMManager
     from lightspeed_evaluation.core.llm.manager import LLMManager
     from lightspeed_evaluation.core.llm.ragas import RagasLLMManager
@@ -19,7 +23,7 @@
     "LLMError": ("lightspeed_evaluation.core.system.exceptions", "LLMError"),
     "LLMManager": ("lightspeed_evaluation.core.llm.manager", "LLMManager"),
     "BaseCustomLLM": ("lightspeed_evaluation.core.llm.custom", "BaseCustomLLM"),
-    "TokenTracker": ("lightspeed_evaluation.core.llm.custom", "TokenTracker"),
+    "TokenTracker": ("lightspeed_evaluation.core.llm.token_tracker", "TokenTracker"),
     "DeepEvalLLMManager": (
         "lightspeed_evaluation.core.llm.deepeval",
         "DeepEvalLLMManager",

src/lightspeed_evaluation/core/llm/custom.py

@@ -1,8 +1,7 @@
 """Base Custom LLM class for evaluation framework."""
 
-import os
 import logging
-import threading
+import os
 from typing import Any, Optional, Union
 
 import litellm
@@ -12,77 +11,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Thread-local storage for active TokenTracker
-_active_tracker: threading.local = threading.local()
-
-
-class TokenTracker:
-    """Tracks token usage from LLM calls using direct response extraction.
-
-    Uses thread-local storage to track the active tracker. Tokens are captured
-    directly from litellm response in BaseCustomLLM.call() - no callbacks,
-    no timeouts, no race conditions.
-
-    Usage:
-        tracker = TokenTracker()
-        tracker.start()  # Set as active tracker for this thread
-        # ... make LLM calls (tokens captured automatically) ...
-        tracker.stop()   # Unset as active tracker
-        input_tokens, output_tokens = tracker.get_counts()
-    """
-
-    def __init__(self) -> None:
-        """Initialize token tracker."""
-        self.input_tokens = 0
-        self.output_tokens = 0
-        self._lock = threading.Lock()  # Instance lock for token counter updates
-
-    def add_tokens(self, prompt_tokens: int, completion_tokens: int) -> None:
-        """Add token counts (thread-safe).
-
-        Called by BaseCustomLLM.call() to record tokens from LLM response.
-
-        Args:
-            prompt_tokens: Number of input/prompt tokens.
-            completion_tokens: Number of output/completion tokens.
-        """
-        with self._lock:
-            self.input_tokens += prompt_tokens
-            self.output_tokens += completion_tokens
-
-    def start(self) -> None:
-        """Set this tracker as active for the current thread."""
-        _active_tracker.tracker = self
-
-    def stop(self) -> None:
-        """Unset this tracker as active for the current thread."""
-        if getattr(_active_tracker, "tracker", None) is self:
-            _active_tracker.tracker = None
-
-    def get_counts(self) -> tuple[int, int]:
-        """Get accumulated token counts.
-
-        Returns:
-            Tuple of (input_tokens, output_tokens)
-        """
-        with self._lock:
-            return self.input_tokens, self.output_tokens
-
-    def reset(self) -> None:
-        """Reset token counts to zero."""
-        with self._lock:
-            self.input_tokens = 0
-            self.output_tokens = 0
-
-    @staticmethod
-    def get_active() -> Optional["TokenTracker"]:
-        """Get the active tracker for the current thread.
-
-        Returns:
-            The active TokenTracker, or None if no tracker is active.
-        """
-        return getattr(_active_tracker, "tracker", None)
-
 
 class BaseCustomLLM:  # pylint: disable=too-few-public-methods
     """Base LLM class with core calling functionality."""
@@ -178,24 +106,3 @@ def call(
 
         except Exception as e:
             raise LLMError(f"LLM call failed: {str(e)}") from e
-
-        finally:
-            # Track tokens even if the call failed - tokens may have been consumed
-            self._track_tokens(response)
-
-    def _track_tokens(self, response: Any) -> None:
-        """Track JudgeLLM tokens if a tracker is active."""
-        # Only track token counts if response exists and is NOT from cache
-        tracker = TokenTracker.get_active()
-        if tracker and response is not None:
-            cache_hit = getattr(
-                response, "_hidden_params", {}
-            ).get(  # pylint: disable=protected-access
-                "cache_hit", False
-            )
-            # Only add tokens if this response was not retrieved from cache
-            if not cache_hit and hasattr(response, "usage") and response.usage:
-                tracker.add_tokens(
-                    getattr(response.usage, "prompt_tokens", 0),
-                    getattr(response.usage, "completion_tokens", 0),
-                )

src/lightspeed_evaluation/core/llm/deepeval.py

@@ -1,4 +1,8 @@
-"""DeepEval LLM Manager - DeepEval-specific LLM wrapper."""
+"""DeepEval LLM Manager - DeepEval-specific LLM wrapper.
+
+Note: litellm patching is applied at package level (__init__.py) before any imports.
+This ensures DeepEval's LiteLLMModel uses the patched completion functions.
+"""
 
 import os
 from typing import Any
@@ -10,7 +14,8 @@
 class DeepEvalLLMManager:
     """DeepEval LLM Manager - Takes LLM parameters directly.
 
-    This manager focuses solely on DeepEval-specific LLM integration.
+    This manager focuses solely on DeepEval-specific LLM integration
+    with token tracking support.
     """
 
     def __init__(self, model_name: str, llm_params: dict[str, Any]):
@@ -23,7 +28,10 @@ def __init__(self, model_name: str, llm_params: dict[str, Any]):
         # Always drop unsupported parameters for cross-provider compatibility
         litellm.drop_params = True
 
-        # Create DeepEval's LLM model with provided parameters
+        # Note: Token tracking is handled by the patched litellm.completion/acompletion
+        # No additional setup needed - the patch was applied at module import time
+
+        # Create standard LiteLLMModel - it will use our patched completion functions
         self.llm_model = LiteLLMModel(
             model=self.model_name,
             temperature=llm_params.get("temperature", 0.0),

src/lightspeed_evaluation/core/llm/litellm_patch.py

@@ -0,0 +1,47 @@
+"""Global litellm patching for token tracking.
+
+It patches litellm.completion and litellm.acompletion to automatically track tokens
+for all LLM calls throughout the application.
+"""
+
+import logging
+from functools import wraps
+from typing import Any
+
+import litellm
+
+from lightspeed_evaluation.core.llm.token_tracker import track_tokens
+
+logger = logging.getLogger(__name__)
+
+
+# Store original functions before patching
+_original_completion = litellm.completion
+_original_acompletion = litellm.acompletion
+
+
+@wraps(_original_completion)
+def _completion_with_token_tracking(*args: Any, **kwargs: Any) -> Any:
+    """Wrapper around litellm.completion that tracks tokens."""
+    response = _original_completion(*args, **kwargs)
+    try:
+        track_tokens(response)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.exception("Failed to track tokens for completion: %s", e)
+    return response
+
+
+@wraps(_original_acompletion)
+async def _acompletion_with_token_tracking(*args: Any, **kwargs: Any) -> Any:
+    """Wrapper around litellm.acompletion that tracks tokens."""
+    response = await _original_acompletion(*args, **kwargs)
+    try:
+        track_tokens(response)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.exception("Failed to track tokens for acompletion: %s", e)
+    return response
+
+
+# Patch litellm's completion functions to include token tracking
+litellm.completion = _completion_with_token_tracking
+litellm.acompletion = _acompletion_with_token_tracking

src/lightspeed_evaluation/core/llm/token_tracker.py

@@ -0,0 +1,102 @@
+"""TokenTracker for tracking LLM token usage with direct response extraction."""
+
+import logging
+import threading
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+# Thread-local storage for active TokenTracker
+_active_tracker: threading.local = threading.local()
+
+
+class TokenTracker:
+    """Tracks token usage from LLM calls using direct response extraction.
+
+    Uses thread-local storage to track the active tracker. Tokens are captured
+    directly from litellm response in BaseCustomLLM.call() - no callbacks,
+    no timeouts, no race conditions.
+
+    Usage:
+        tracker = TokenTracker()
+        tracker.start()  # Set as active tracker for this thread
+        # ... make LLM calls (tokens captured automatically) ...
+        tracker.stop()   # Unset as active tracker
+        input_tokens, output_tokens = tracker.get_counts()
+    """
+
+    def __init__(self) -> None:
+        """Initialize token tracker."""
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self._lock = threading.Lock()  # Instance lock for token counter updates
+
+    def add_tokens(self, prompt_tokens: int, completion_tokens: int) -> None:
+        """Add token counts (thread-safe).
+
+        Called by BaseCustomLLM.call() to record tokens from LLM response.
+
+        Args:
+            prompt_tokens: Number of input/prompt tokens.
+            completion_tokens: Number of output/completion tokens.
+        """
+        with self._lock:
+            self.input_tokens += prompt_tokens
+            self.output_tokens += completion_tokens
+
+    def start(self) -> None:
+        """Set this tracker as active for the current thread."""
+        _active_tracker.tracker = self
+
+    def stop(self) -> None:
+        """Unset this tracker as active for the current thread."""
+        if getattr(_active_tracker, "tracker", None) is self:
+            _active_tracker.tracker = None
+
+    def get_counts(self) -> tuple[int, int]:
+        """Get accumulated token counts.
+
+        Returns:
+            Tuple of (input_tokens, output_tokens)
+        """
+        with self._lock:
+            return self.input_tokens, self.output_tokens
+
+    def reset(self) -> None:
+        """Reset token counts to zero."""
+        with self._lock:
+            self.input_tokens = 0
+            self.output_tokens = 0
+
+    @staticmethod
+    def get_active() -> Optional["TokenTracker"]:
+        """Get the active tracker for the current thread.
+
+        Returns:
+            The active TokenTracker, or None if no tracker is active.
+        """
+        return getattr(_active_tracker, "tracker", None)
+
+
+def track_tokens(response: Any) -> None:
+    """Track JudgeLLM tokens if a tracker is active.
+
+    Called by the litellm patch (see llm_patch.py) after each LLM call.
+    Skips tracking for cached responses to avoid counting tokens that weren't actually consumed.
+    """
+    # Only track token counts if response exists and is NOT from cache
+    tracker = TokenTracker.get_active()
+    if tracker and response is not None:
+        cache_hit = getattr(
+            response, "_hidden_params", {}
+        ).get(  # pylint: disable=protected-access
+            "cache_hit", False
+        )
+        # Only add tokens if this response was not retrieved from cache
+        if not cache_hit and hasattr(response, "usage") and response.usage:
+            prompt_tokens = int(getattr(response.usage, "prompt_tokens", 0))
+            completion_tokens = int(getattr(response.usage, "completion_tokens", 0))
+            tracker.add_tokens(
+                prompt_tokens,
+                completion_tokens,
+            )

src/lightspeed_evaluation/pipeline/evaluation/evaluator.py

@@ -6,7 +6,7 @@
 from typing import Any, Optional
 
 from lightspeed_evaluation.core.embedding.manager import EmbeddingManager
-from lightspeed_evaluation.core.llm.custom import TokenTracker
+from lightspeed_evaluation.core.llm.token_tracker import TokenTracker
 from lightspeed_evaluation.core.llm.manager import LLMManager
 from lightspeed_evaluation.core.metrics.custom import CustomMetrics
 from lightspeed_evaluation.core.metrics.deepeval import DeepEvalMetrics