scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

refactor: move callback failure tracking to wrapper level

Browse source 
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
This commit is contained in:
Sameer Kankute 2026-03-10 18:00:01 +05:30
parent 11b4c456b4
commit 0ee1da9b95
4 changed files with 200 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

92
litellm/integrations/langfuse/langfuse_prompt_management.py
View file

@ -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,
)

3
litellm/integrations/opentelemetry.py
View file

@ -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)

4
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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__"):

207
tests/enterprise/litellm_enterprise/integrations/test_prometheus.py
View file

@ -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
# ==============================================================================