From 0ee1da9b95089ef9d339c2c3ef2183aa8f6398f1 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 10 Mar 2026 18:00:01 +0530 Subject: [PATCH] refactor: move callback failure tracking to wrapper level Move litellm_callback_logging_failures_metric tracking from individual callbacks to the wrapper's exception handler in litellm_logging.py. This provides: - Single source of truth for all callback failures - Automatic coverage for all callbacks (langfuse, langfuse_otel, otel, s3, etc.) - No need for custom handling in each callback implementation Changes: - Enhanced _get_callback_name() to check callback_name and integration_name attributes - Removed try-catch and handle_callback_failure() from langfuse_prompt_management.py - Added raise to opentelemetry.py exception handlers to bubble up to wrapper - Updated tests to verify wrapper-level tracking for langfuse, langfuse_otel, and otel Made-with: Cursor --- .../langfuse/langfuse_prompt_management.py | 92 ++++---- litellm/integrations/opentelemetry.py | 3 +- litellm/litellm_core_utils/litellm_logging.py | 4 + .../integrations/test_prometheus.py | 207 +++++++++++++----- 4 files changed, 200 insertions(+), 106 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 3986fc6a6ef..8f73eabad44 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -300,59 +300,43 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - standard_callback_dynamic_params = kwargs.get( - "standard_callback_dynamic_params" - ) - langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=self, - standard_callback_dynamic_params=standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) - langfuse_logger_to_use.log_event_on_langfuse( - kwargs=kwargs, - response_obj=response_obj, - start_time=start_time, - end_time=end_time, - user_id=kwargs.get("user", None), - ) - except Exception as e: - from litellm._logging import verbose_logger - - verbose_logger.exception( - f"Langfuse Layer Error - Exception occurred while logging success event: {str(e)}" - ) - self.handle_callback_failure(callback_name="langfuse") + standard_callback_dynamic_params = kwargs.get( + "standard_callback_dynamic_params" + ) + langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( + globalLangfuseLogger=self, + standard_callback_dynamic_params=standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + langfuse_logger_to_use.log_event_on_langfuse( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + user_id=kwargs.get("user", None), + ) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - standard_callback_dynamic_params = kwargs.get( - "standard_callback_dynamic_params" - ) - langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=self, - standard_callback_dynamic_params=standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) - standard_logging_object = cast( - Optional[StandardLoggingPayload], - kwargs.get("standard_logging_object", None), - ) - if standard_logging_object is None: - return - langfuse_logger_to_use.log_event_on_langfuse( - start_time=start_time, - end_time=end_time, - response_obj=None, - user_id=kwargs.get("user", None), - status_message=standard_logging_object["error_str"], - level="ERROR", - kwargs=kwargs, - ) - except Exception as e: - from litellm._logging import verbose_logger - - verbose_logger.exception( - f"Langfuse Layer Error - Exception occurred while logging failure event: {str(e)}" - ) - self.handle_callback_failure(callback_name="langfuse") + standard_callback_dynamic_params = kwargs.get( + "standard_callback_dynamic_params" + ) + langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( + globalLangfuseLogger=self, + standard_callback_dynamic_params=standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + standard_logging_object = cast( + Optional[StandardLoggingPayload], + kwargs.get("standard_logging_object", None), + ) + if standard_logging_object is None: + return + langfuse_logger_to_use.log_event_on_langfuse( + start_time=start_time, + end_time=end_time, + response_obj=None, + user_id=kwargs.get("user", None), + status_message=standard_logging_object["error_str"], + level="ERROR", + kwargs=kwargs, + ) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index dbc1cb0d6b0..3476f8a2f8a 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1609,10 +1609,10 @@ class OpenTelemetry(CustomLogger): ) except Exception as e: - self.handle_callback_failure(callback_name= self.callback_name) verbose_logger.exception( "OpenTelemetry logging error in set_attributes %s", str(e) ) + raise def _cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: """ @@ -1749,6 +1749,7 @@ class OpenTelemetry(CustomLogger): verbose_logger.exception( "OpenTelemetry logging error in set_raw_request_attributes %s", str(e) ) + raise def _to_ns(self, dt): return int(dt.timestamp() * 1e9) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d3bcfe8200e..805939fabae 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3092,6 +3092,10 @@ class Logging(LiteLLMLoggingBaseClass): """ if isinstance(cb, str): return cb + if hasattr(cb, "callback_name") and cb.callback_name: + return cb.callback_name + if hasattr(cb, "integration_name"): + return cb.integration_name if hasattr(cb, "__name__"): return cb.__name__ if hasattr(cb, "__func__"): diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py index 4a08c208fda..1b59c1d4da5 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py @@ -942,14 +942,17 @@ async def test_langfuse_callback_failure_metric(prometheus_logger): """ Test that Langfuse callback failures are properly tracked in Prometheus metrics. - This test verifies that when Langfuse logging fails, the - litellm_callback_logging_failures_metric is incremented with callback_name="langfuse". + This test verifies that when Langfuse logging fails, the wrapper-level exception + handler in litellm_logging.py increments litellm_callback_logging_failures_metric + with callback_name="langfuse". """ + from datetime import datetime from unittest.mock import MagicMock, patch from litellm.integrations.langfuse.langfuse_prompt_management import ( LangfusePromptManagement, ) + from litellm.litellm_core_utils.litellm_logging import Logging # Get initial value initial_value = 0 @@ -964,6 +967,20 @@ async def test_langfuse_callback_failure_metric(prometheus_logger): with patch("litellm.integrations.langfuse.langfuse_prompt_management.langfuse_client_init"): langfuse_logger = LangfusePromptManagement() + # Register prometheus logger in litellm callbacks + litellm.callbacks = [prometheus_logger] + + # Create a Logging wrapper instance + logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + # Mock the log_event_on_langfuse to raise an exception with patch( "litellm.integrations.langfuse.langfuse_prompt_management.LangFuseHandler.get_langfuse_logger_for_request" @@ -972,23 +989,28 @@ async def test_langfuse_callback_failure_metric(prometheus_logger): mock_logger.log_event_on_langfuse.side_effect = Exception("Langfuse API error") mock_get_logger.return_value = mock_logger - # Mock handle_callback_failure to track calls - with patch.object(prometheus_logger, "increment_callback_logging_failure") as mock_increment: - # Inject prometheus logger into the langfuse logger - langfuse_logger.handle_callback_failure = lambda callback_name: mock_increment( - callback_name=callback_name + # Call the wrapper's async_success_handler with langfuse callback + # The wrapper should catch the exception and call _handle_callback_failure + logging_obj.dynamic_async_success_callbacks = [langfuse_logger] + + try: + await logging_obj.async_success_handler( + result={"choices": [{"message": {"role": "assistant", "content": "test"}}]}, + start_time=datetime.now(), + end_time=datetime.now(), ) - - # Call async_log_success_event - should catch exception and increment metric - await langfuse_logger.async_log_success_event( - kwargs={}, - response_obj={}, - start_time=None, - end_time=None, - ) - - # Verify that increment was called with correct callback name - mock_increment.assert_called_once_with(callback_name="langfuse") + except Exception: + pass + + # Verify that the metric was incremented + current_value = prometheus_logger.litellm_callback_logging_failures_metric.labels( + callback_name="langfuse" + )._value.get() + + assert current_value == initial_value + 1, ( + f"Expected callback failure metric to increment by 1, " + f"got {current_value - initial_value}" + ) print("✓ Langfuse callback failure metric test passed") @@ -998,12 +1020,15 @@ async def test_langfuse_otel_callback_failure_metric(prometheus_logger): """ Test that Langfuse OTEL callback failures are properly tracked in Prometheus metrics. - This test verifies that when Langfuse OTEL logging fails, the - litellm_callback_logging_failures_metric is incremented with callback_name="langfuse_otel". + This test verifies that when Langfuse OTEL logging fails, the wrapper-level exception + handler in litellm_logging.py increments litellm_callback_logging_failures_metric + with callback_name="langfuse_otel". """ + from datetime import datetime from unittest.mock import MagicMock, patch from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger + from litellm.litellm_core_utils.litellm_logging import Logging # Get initial value initial_value = 0 @@ -1019,42 +1044,122 @@ async def test_langfuse_otel_callback_failure_metric(prometheus_logger): langfuse_otel_logger = LangfuseOtelLogger(callback_name="langfuse_otel") langfuse_otel_logger.callback_name = "langfuse_otel" - # Mock handle_callback_failure to track calls - with patch.object(prometheus_logger, "increment_callback_logging_failure") as mock_increment: - # Inject prometheus logger into the langfuse otel logger - langfuse_otel_logger.handle_callback_failure = lambda callback_name: mock_increment( - callback_name=callback_name - ) + # Register prometheus logger in litellm callbacks + litellm.callbacks = [prometheus_logger] + + # Create a Logging wrapper instance + logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + + # Mock set_attributes to raise an exception + with patch.object(langfuse_otel_logger, "set_attributes") as mock_set_attributes: + mock_set_attributes.side_effect = Exception("OTEL attribute error") - # Test that the OpenTelemetry base class set_attributes exception handler works - # This is where langfuse_otel failures are caught and tracked - with patch.object(langfuse_otel_logger, "set_attributes") as mock_set_attributes: - # Simulate the exception handling in set_attributes - def set_attributes_with_error(*args, **kwargs): - # This simulates what happens in the real set_attributes method - try: - raise Exception("Attribute error") - except Exception as e: - langfuse_otel_logger.handle_callback_failure(callback_name=langfuse_otel_logger.callback_name) - - mock_set_attributes.side_effect = set_attributes_with_error - - # Call set_attributes - try: - langfuse_otel_logger.set_attributes( - span=MagicMock(), - kwargs={}, - response_obj={} - ) - except Exception: - pass - - # Verify that increment was called with correct callback name - mock_increment.assert_called_with(callback_name="langfuse_otel") + # Call the wrapper's async_success_handler with langfuse_otel callback + # The wrapper should catch the exception and call _handle_callback_failure + logging_obj.dynamic_async_success_callbacks = [langfuse_otel_logger] + + try: + await logging_obj.async_success_handler( + result={"choices": [{"message": {"role": "assistant", "content": "test"}}]}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + except Exception: + pass + + # Verify that the metric was incremented + current_value = prometheus_logger.litellm_callback_logging_failures_metric.labels( + callback_name="langfuse_otel" + )._value.get() + + assert current_value == initial_value + 1, ( + f"Expected callback failure metric to increment by 1, " + f"got {current_value - initial_value}" + ) print("✓ Langfuse OTEL callback failure metric test passed") +@pytest.mark.asyncio +async def test_generic_otel_callback_failure_metric(prometheus_logger): + """ + Test that generic OTEL callback failures are properly tracked in Prometheus metrics. + + This test verifies that the wrapper-level exception handler works for any OTEL-based + callback, not just langfuse_otel. + """ + from datetime import datetime + from unittest.mock import MagicMock, patch + + from litellm.integrations.opentelemetry import OpenTelemetry + from litellm.litellm_core_utils.litellm_logging import Logging + + # Get initial value + initial_value = 0 + try: + initial_value = prometheus_logger.litellm_callback_logging_failures_metric.labels( + callback_name="otel" + )._value.get() + except Exception: + initial_value = 0 + + # Create generic OTEL logger with mocked initialization + with patch("litellm.integrations.opentelemetry.OpenTelemetry.__init__", return_value=None): + otel_logger = OpenTelemetry(callback_name="otel") + otel_logger.callback_name = "otel" + + # Register prometheus logger in litellm callbacks + litellm.callbacks = [prometheus_logger] + + # Create a Logging wrapper instance + logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + + # Mock set_attributes to raise an exception + with patch.object(otel_logger, "set_attributes") as mock_set_attributes: + mock_set_attributes.side_effect = Exception("OTEL attribute error") + + # Call the wrapper's async_success_handler with otel callback + # The wrapper should catch the exception and call _handle_callback_failure + logging_obj.dynamic_async_success_callbacks = [otel_logger] + + try: + await logging_obj.async_success_handler( + result={"choices": [{"message": {"role": "assistant", "content": "test"}}]}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + except Exception: + pass + + # Verify that the metric was incremented + current_value = prometheus_logger.litellm_callback_logging_failures_metric.labels( + callback_name="otel" + )._value.get() + + assert current_value == initial_value + 1, ( + f"Expected callback failure metric to increment by 1, " + f"got {current_value - initial_value}" + ) + + print("✓ Generic OTEL callback failure metric test passed") + + # ============================================================================== # END CALLBACK FAILURE METRICS TESTS # ==============================================================================