scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

Merge pull request #24572 from Sameerlite/litellm_callback_failure_metrics_clean

Browse source 
fix(callbacks): unify callback failure metrics across integrations
This commit is contained in:
Sameer Kankute 2026-04-02 18:22:23 +05:30 • committed by GitHub
commit c8cb49fdfb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 321 additions and 170 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/integrations/custom_batch_logger.py
View file

@ -30,6 +30,8 @@ class CustomBatchLogger(CustomLogger):
self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE
self.last_flush_time = time.time()
self.flush_lock = flush_lock
# Track health transition so we emit one metric when callback goes unhealthy.
self._is_in_failure_state: bool = False
super().__init__(**kwargs)
@ -50,9 +52,33 @@ class CustomBatchLogger(CustomLogger):
verbose_logger.debug(
"CustomLogger: Flushing batch of %s events", len(self.log_queue)
)
await self.async_send_batch()
try:
await self.async_send_batch()
except Exception as e:
callback_name = self._get_callback_failure_name()
verbose_logger.debug(
"CustomBatchLogger: batch flush failed for %s: %s",
callback_name,
str(e),
)
if not self._is_in_failure_state:
self._report_callback_failure(callback_name=callback_name)
self._is_in_failure_state = True
else:
verbose_logger.debug(
"CustomBatchLogger: callback %s already in failure state, skipping duplicate metric",
callback_name,
)
return
self.log_queue.clear()
self.last_flush_time = time.time()
self._is_in_failure_state = False
def _get_callback_failure_name(self) -> str:
"""
Default callback name for batch failures.
"""
return self.__class__.__name__
async def async_send_batch(self, *args, **kwargs):
pass

19
litellm/integrations/custom_logger.py
View file

@ -795,9 +795,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
This is useful for logging payloads that contain sensitive information.
"""
import litellm
from copy import copy
import litellm
from litellm import Choices, Message, ModelResponse
turn_off_message_logging: bool = getattr(
@ -888,11 +888,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
pass
def handle_callback_failure(self, callback_name: str):
def _report_callback_failure(self, callback_name: str) -> None:
"""
Handle callback logging failures by incrementing Prometheus metrics.
Report callback failures to Prometheus callback failure metric.
Call this method in exception handlers within your callback when logging fails.
Keep this as a single helper so callback integrations have one
standard failure-reporting path.
"""
try:
import litellm
@ -917,9 +918,17 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
from litellm._logging import verbose_logger
verbose_logger.debug(
f"Error in handle_callback_failure for {callback_name}: {str(e)}"
f"Error in _report_callback_failure for {callback_name}: {str(e)}"
)
def handle_callback_failure(self, callback_name: str):
"""
Handle callback logging failures by incrementing Prometheus metrics.
Call this method in exception handlers within your callback when logging fails.
"""
self._report_callback_failure(callback_name=callback_name)
async def _strip_base64_from_messages(
self,
payload: "StandardLoggingPayload",

124
litellm/integrations/datadog/datadog.py
View file

@ -27,16 +27,16 @@ import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_mock_client import (
should_use_datadog_mock,
create_mock_datadog_client,
)
from litellm.integrations.datadog.datadog_handler import (
get_datadog_base_url_from_env,
get_datadog_hostname,
get_datadog_service,
get_datadog_source,
get_datadog_tags,
get_datadog_base_url_from_env,
)
from litellm.integrations.datadog.datadog_mock_client import (
create_mock_datadog_client,
should_use_datadog_mock,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.llms.custom_httpx.http_handler import (
@ -48,10 +48,10 @@ from litellm.types.integrations.base_health_check import IntegrationHealthCheckS
from litellm.types.integrations.datadog import (
DD_ERRORS,
DD_MAX_BATCH_SIZE,
DataDogStatus,
DatadogInitParams,
DatadogPayload,
DatadogProxyFailureHookJsonMessage,
DataDogStatus,
)
from litellm.types.services import ServiceLoggerPayload, ServiceTypes
from litellm.types.utils import StandardLoggingPayload
@ -191,32 +191,18 @@ class DataDogLogger(
Raises:
Raises a NON Blocking verbose_logger.exception if an error occurs
Raises exceptions to framework wrapper which handles metric tracking
"""
try:
verbose_logger.debug(
"Datadog: Logging - Enters logging function for model %s", kwargs
)
await self._log_async_event(kwargs, response_obj, start_time, end_time)
except Exception as e:
verbose_logger.exception(
f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}"
)
pass
verbose_logger.debug(
"Datadog: Logging - Enters logging function for model %s", kwargs
)
await self._log_async_event(kwargs, response_obj, start_time, end_time)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
"Datadog: Logging - Enters logging function for model %s", kwargs
)
await self._log_async_event(kwargs, response_obj, start_time, end_time)
except Exception as e:
verbose_logger.exception(
f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}"
)
pass
verbose_logger.debug(
"Datadog: Logging - Enters logging function for model %s", kwargs
)
await self._log_async_event(kwargs, response_obj, start_time, end_time)
async def async_post_call_failure_hook(
self,
@ -301,11 +287,12 @@ class DataDogLogger(
self.log_queue.append(dd_payload)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
await self.flush_queue()
except Exception as e:
verbose_logger.exception(
f"Datadog: async_post_call_failure_hook - {str(e)}\n{traceback.format_exc()}"
)
raise
return None
async def async_send_batch(self):
@ -360,6 +347,7 @@ class DataDogLogger(
verbose_logger.exception(
f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}"
)
raise
def log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
@ -368,52 +356,44 @@ class DataDogLogger(
- Creates a Datadog payload
- instantly logs it on DD API
"""
try:
if litellm.datadog_use_v1 is True:
dd_payload = self._create_v0_logging_payload(
kwargs=kwargs,
response_obj=response_obj,
start_time=start_time,
end_time=end_time,
)
else:
dd_payload = self.create_datadog_logging_payload(
kwargs=kwargs,
response_obj=response_obj,
start_time=start_time,
end_time=end_time,
)
# Build headers
headers = {}
# Add API key if available (required for direct API, optional for agent)
if self.DD_API_KEY:
headers["DD-API-KEY"] = self.DD_API_KEY
response = self.sync_client.post(
url=self.intake_url,
json=dd_payload, # type: ignore
headers=headers,
if litellm.datadog_use_v1 is True:
dd_payload = self._create_v0_logging_payload(
kwargs=kwargs,
response_obj=response_obj,
start_time=start_time,
end_time=end_time,
)
else:
dd_payload = self.create_datadog_logging_payload(
kwargs=kwargs,
response_obj=response_obj,
start_time=start_time,
end_time=end_time,
)
response.raise_for_status()
if response.status_code != 202:
raise Exception(
f"Response from datadog API status_code: {response.status_code}, text: {response.text}"
)
# Build headers
headers = {}
# Add API key if available (required for direct API, optional for agent)
if self.DD_API_KEY:
headers["DD-API-KEY"] = self.DD_API_KEY
verbose_logger.debug(
"Datadog: Response from datadog API status_code: %s, text: %s",
response.status_code,
response.text,
response = self.sync_client.post(
url=self.intake_url,
json=dd_payload, # type: ignore
headers=headers,
)
response.raise_for_status()
if response.status_code != 202:
raise Exception(
f"Response from datadog API status_code: {response.status_code}, text: {response.text}"
)
except Exception as e:
verbose_logger.exception(
f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}"
)
pass
pass
verbose_logger.debug(
"Datadog: Response from datadog API status_code: %s, text: %s",
response.status_code,
response.text,
)
async def _log_async_event(self, kwargs, response_obj, start_time, end_time):
dd_payload = self.create_datadog_logging_payload(
@ -429,7 +409,7 @@ class DataDogLogger(
)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
await self.flush_queue()
def _create_datadog_logging_payload_helper(
self,

56
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -9,7 +9,6 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp
import asyncio
import json
import os
from litellm._uuid import uuid
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Union
@ -17,15 +16,16 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.integrations.datadog.datadog_mock_client import (
should_use_datadog_mock,
create_mock_datadog_client,
)
from litellm.integrations.datadog.datadog_handler import (
get_datadog_base_url_from_env,
get_datadog_service,
get_datadog_tags,
get_datadog_base_url_from_env,
)
from litellm.integrations.datadog.datadog_mock_client import (
create_mock_datadog_client,
should_use_datadog_mock,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -152,36 +152,26 @@ class DataDogLLMObsLogger(CustomBatchLogger):
return dict_datadog_llm_obs_params
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}"
)
payload = self.create_llm_obs_payload(kwargs, start_time, end_time)
verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}")
self.log_queue.append(payload)
verbose_logger.debug(
f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}"
)
payload = self.create_llm_obs_payload(kwargs, start_time, end_time)
verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}")
self.log_queue.append(payload)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
except Exception as e:
verbose_logger.exception(
f"DataDogLLMObs: Error logging success event - {str(e)}"
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
f"DataDogLLMObs: Logging failure event for model {kwargs.get('model', 'unknown')}"
)
payload = self.create_llm_obs_payload(kwargs, start_time, end_time)
verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}")
self.log_queue.append(payload)
verbose_logger.debug(
f"DataDogLLMObs: Logging failure event for model {kwargs.get('model', 'unknown')}"
)
payload = self.create_llm_obs_payload(kwargs, start_time, end_time)
verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}")
self.log_queue.append(payload)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
except Exception as e:
verbose_logger.exception(
f"DataDogLLMObs: Error logging failure event - {str(e)}"
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def async_send_batch(self):
try:
@ -249,8 +239,10 @@ class DataDogLLMObsLogger(CustomBatchLogger):
verbose_logger.exception(
f"DataDogLLMObs: Error sending batch - {e.response.text}"
)
raise
except Exception as e:
verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {str(e)}")
raise
def create_llm_obs_payload(
self, kwargs: Dict, start_time: datetime, end_time: datetime

66
litellm/integrations/datadog/datadog_metrics.py
View file

@ -155,55 +155,41 @@ class DatadogMetricsLogger(CustomBatchLogger):
self.log_queue.append(series_count)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_object is None:
return
if standard_logging_object is None:
return
self._add_metrics_from_log(
log=standard_logging_object, kwargs=kwargs, status_code="200"
)
self._add_metrics_from_log(
log=standard_logging_object, kwargs=kwargs, status_code="200"
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
except Exception as e:
verbose_logger.exception(
f"Datadog Metrics: Error in async_log_success_event: {str(e)}"
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
try:
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_object is None:
return
if standard_logging_object is None:
return
# Extract status code from error information
status_code = "500" # default
error_information = (
standard_logging_object.get("error_information", {}) or {}
)
error_code = error_information.get("error_code") # type: ignore
if error_code is not None:
status_code = str(error_code)
# Extract status code from error information
status_code = "500" # default
error_information = standard_logging_object.get("error_information", {}) or {}
error_code = error_information.get("error_code") # type: ignore
if error_code is not None:
status_code = str(error_code)
self._add_metrics_from_log(
log=standard_logging_object, kwargs=kwargs, status_code=status_code
)
self._add_metrics_from_log(
log=standard_logging_object, kwargs=kwargs, status_code=status_code
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
except Exception as e:
verbose_logger.exception(
f"Datadog Metrics: Error in async_log_failure_event: {str(e)}"
)
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def async_send_batch(self):
if not self.log_queue:

102
litellm/integrations/opentelemetry.py
View file

@ -1,4 +1,5 @@
import os
import types
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
@ -160,6 +161,7 @@ class OpenTelemetry(CustomLogger):
self.OTEL_EXPORTER = self.config.exporter
self.OTEL_ENDPOINT = self.config.endpoint
self.OTEL_HEADERS = self.config.headers
self._span_export_in_failure_state = False
self._tracer_provider_cache: Dict[str, Any] = {}
self._init_tracing(tracer_provider)
@ -294,7 +296,7 @@ class OpenTelemetry(CustomLogger):
return provider
def _init_tracing(self, tracer_provider):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import SpanKind
@ -333,7 +335,7 @@ class OpenTelemetry(CustomLogger):
self._response_duration_histogram = None
return
from opentelemetry import metrics
import opentelemetry.metrics as metrics
from opentelemetry.sdk.metrics import MeterProvider
def create_meter_provider():
@ -433,7 +435,7 @@ class OpenTelemetry(CustomLogger):
end_time: Optional[Union[datetime, float]] = None,
event_metadata: Optional[dict] = None,
):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.trace import Status, StatusCode
_start_time_ns = 0
@ -493,7 +495,7 @@ class OpenTelemetry(CustomLogger):
end_time: Optional[Union[float, datetime]] = None,
event_metadata: Optional[dict] = None,
):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.trace import Status, StatusCode
_start_time_ns = 0
@ -555,7 +557,7 @@ class OpenTelemetry(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
traceback_str: Optional[str] = None,
):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.trace import Status, StatusCode
parent_otel_span = user_api_key_dict.parent_otel_span
@ -746,7 +748,7 @@ class OpenTelemetry(CustomLogger):
span = None
# Only set attributes if the span is still recording (not closed)
# Note: parent_span is guaranteed to be not None here
if hasattr(parent_span, "set_status"):
if parent_span is not None and hasattr(parent_span, "set_status"):
parent_span.set_status(Status(StatusCode.OK))
self.set_attributes(parent_span, kwargs, response_obj)
# Raw-request as direct child of parent_span
@ -808,7 +810,7 @@ class OpenTelemetry(CustomLogger):
def _maybe_log_raw_request(
self, kwargs, response_obj, start_time, end_time, parent_span
):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.trace import Status, StatusCode
# only log raw LLM request/response if message_logging is on and not globally turned off
@ -1076,8 +1078,8 @@ class OpenTelemetry(CustomLogger):
from opentelemetry._logs import SeverityNumber, get_logger
try:
from opentelemetry.sdk._logs import ( # type: ignore[attr-defined] # OTEL < 1.39.0
LogRecord as SdkLogRecord,
from opentelemetry.sdk._logs import (
LogRecord as SdkLogRecord, # type: ignore[attr-defined] # OTEL < 1.39.0
)
except ImportError:
from opentelemetry.sdk._logs._internal import (
@ -1159,7 +1161,7 @@ class OpenTelemetry(CustomLogger):
2. The parent proxy-request span
3. The original fallback context (may be None — last resort)
"""
from opentelemetry import trace as _trace
import opentelemetry.trace as _trace
if span is not None:
return _trace.set_span_in_context(span)
@ -1288,7 +1290,7 @@ class OpenTelemetry(CustomLogger):
# record error on parent span (keeps hierarchy shallow)
# Only set attributes if the span is still recording (not closed)
# Note: parent_otel_span is guaranteed to be not None here
if parent_otel_span.is_recording():
if parent_otel_span is not None and parent_otel_span.is_recording():
parent_otel_span.set_status(Status(StatusCode.ERROR))
self.set_attributes(parent_otel_span, kwargs, response_obj)
self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs)
@ -1923,7 +1925,8 @@ class OpenTelemetry(CustomLogger):
return _parent_context
def _get_span_context(self, kwargs, default_span: Optional[Span] = None):
from opentelemetry import context, trace
import opentelemetry.context as context
import opentelemetry.trace as trace
from opentelemetry.trace.propagation.tracecontext import (
TraceContextTextMapPropagator,
)
@ -2016,7 +2019,10 @@ class OpenTelemetry(CustomLogger):
"OpenTelemetry: intiializing SpanExporter. Value of OTEL_EXPORTER: %s",
self.OTEL_EXPORTER,
)
return SimpleSpanProcessor(cast(SpanExporter, self.OTEL_EXPORTER))
wrapped_exporter = self._wrap_exporter_with_failure_tracking(
cast(SpanExporter, self.OTEL_EXPORTER)
)
return SimpleSpanProcessor(cast(SpanExporter, wrapped_exporter))
if self.OTEL_EXPORTER == "console":
verbose_logger.debug(
@ -2046,11 +2052,12 @@ class OpenTelemetry(CustomLogger):
normalized_endpoint = self._normalize_otel_endpoint(
self.OTEL_ENDPOINT, "traces"
)
return BatchSpanProcessor(
wrapped_exporter = self._wrap_exporter_with_failure_tracking(
OTLPSpanExporterHTTP(
endpoint=normalized_endpoint, headers=_split_otel_headers
),
)
)
return BatchSpanProcessor(cast(SpanExporter, wrapped_exporter))
elif self.OTEL_EXPORTER == "otlp_grpc" or self.OTEL_EXPORTER == "grpc":
try:
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
@ -2069,11 +2076,12 @@ class OpenTelemetry(CustomLogger):
normalized_endpoint = self._normalize_otel_endpoint(
self.OTEL_ENDPOINT, "traces"
)
return BatchSpanProcessor(
wrapped_exporter = self._wrap_exporter_with_failure_tracking(
OTLPSpanExporterGRPC(
endpoint=normalized_endpoint, headers=_split_otel_headers
),
)
)
return BatchSpanProcessor(cast(SpanExporter, wrapped_exporter))
else:
verbose_logger.debug(
"OpenTelemetry: intiializing console exporter. Value of OTEL_EXPORTER: %s",
@ -2081,6 +2089,62 @@ class OpenTelemetry(CustomLogger):
)
return BatchSpanProcessor(ConsoleSpanExporter())
def _wrap_exporter_with_failure_tracking(self, exporter: SpanExporter) -> Any:
if getattr(exporter, "_litellm_failure_tracking_wrapped", False):
return exporter
original_export = exporter.export
parent = self
def _export_with_failure_tracking(self, spans):
try:
result = original_export(spans)
except Exception:
parent._report_span_export_failure()
raise
if parent._is_failed_span_export_result(result):
parent._report_span_export_failure()
else:
parent._mark_span_export_success()
return result
setattr(
exporter,
"export",
types.MethodType(_export_with_failure_tracking, exporter),
)
setattr(exporter, "_litellm_failure_tracking_wrapped", True)
return exporter
@staticmethod
def _is_failed_span_export_result(result: Any) -> bool:
try:
from opentelemetry.sdk.trace.export import SpanExportResult
return result == SpanExportResult.FAILURE
except Exception:
normalized_result = str(result).strip().lower()
return normalized_result == "failure" or normalized_result.endswith(
".failure"
)
def _get_callback_failure_metric_name(self) -> str:
if self.callback_name:
return self.callback_name
return self.__class__.__name__
def _report_span_export_failure(self) -> None:
if self._span_export_in_failure_state:
return
self._report_callback_failure(
callback_name=self._get_callback_failure_metric_name()
)
self._span_export_in_failure_state = True
def _mark_span_export_success(self) -> None:
self._span_export_in_failure_state = False
def _get_log_exporter(self):
"""
Get the appropriate log exporter based on the configuration.
@ -2327,7 +2391,7 @@ class OpenTelemetry(CustomLogger):
logging_payload: ManagementEndpointLoggingPayload,
parent_otel_span: Optional[Span] = None,
):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.trace import Status, StatusCode
_start_time_ns = 0
@ -2380,7 +2444,7 @@ class OpenTelemetry(CustomLogger):
logging_payload: ManagementEndpointLoggingPayload,
parent_otel_span: Optional[Span] = None,
):
from opentelemetry import trace
import opentelemetry.trace as trace
from opentelemetry.trace import Status, StatusCode
_start_time_ns = 0

5
litellm/litellm_core_utils/litellm_logging.py
View file

@ -3094,6 +3094,11 @@ class Logging(LiteLLMLoggingBaseClass):
)
if capture_exception: # log this error to sentry for debugging
capture_exception(e)
# Track callback logging failures in Prometheus
try:
self._handle_callback_failure(callback=callback)
except Exception:
pass
except Exception as e:
verbose_logger.exception(
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging {}".format(

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

@ -2,6 +2,7 @@
Mock prometheus unit tests, these don't rely on LLM API calls
"""
import asyncio
import json
import os
import sys
@ -13,7 +14,7 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import pytest_asyncio
from apscheduler.schedulers.asyncio import AsyncIOScheduler
@ -1058,6 +1059,94 @@ async def test_langfuse_otel_callback_failure_metric(prometheus_logger):
print("✓ Langfuse OTEL callback failure metric test passed")
@pytest.mark.asyncio
async def test_custom_batch_logger_auto_failure_reporting():
"""
Ensure batch logger failures are reported automatically via shared flush path.
"""
from litellm.integrations.custom_batch_logger import CustomBatchLogger
class _TestBatchLogger(CustomBatchLogger):
def _get_callback_failure_name(self) -> str:
return "AutoTestLogger"
async def async_send_batch(self, *args, **kwargs):
raise RuntimeError("batch send failed")
logger = _TestBatchLogger(flush_lock=asyncio.Lock(), batch_size=1)
logger.log_queue.append({"dummy": "event"})
with patch.object(logger, "_report_callback_failure") as mock_report:
await logger.flush_queue()
mock_report.assert_called_once_with(callback_name="AutoTestLogger")
@pytest.mark.asyncio
async def test_custom_batch_logger_dedupes_repeated_failed_batch_metric():
"""
Ensure periodic retries for the same failed batch do not over-increment metrics.
"""
from litellm.integrations.custom_batch_logger import CustomBatchLogger
class _TestBatchLogger(CustomBatchLogger):
def _get_callback_failure_name(self) -> str:
return "AutoTestLogger"
async def async_send_batch(self, *args, **kwargs):
raise RuntimeError("batch send failed")
logger = _TestBatchLogger(flush_lock=asyncio.Lock(), batch_size=1)
logger.log_queue.append({"dummy": "event"})
with patch.object(logger, "_report_callback_failure") as mock_report:
await logger.flush_queue()
logger.log_queue.append({"dummy": "event-2"})
await logger.flush_queue()
await logger.flush_queue()
mock_report.assert_called_once_with(callback_name="AutoTestLogger")
def test_otel_exporter_failure_reports_callback_metric_once_per_outage():
"""
Ensure OTEL exporter failures increment callback failure metric and dedupe
repeated failures until export recovers.
"""
from opentelemetry.sdk.trace.export import SpanExportResult
from litellm.integrations.opentelemetry import OpenTelemetry
class _DummyExporter:
def __init__(self):
self._results = [
SpanExportResult.FAILURE,
SpanExportResult.FAILURE,
SpanExportResult.SUCCESS,
SpanExportResult.FAILURE,
]
self._idx = 0
def export(self, spans):
result = self._results[self._idx]
self._idx += 1
return result
otel_logger = OpenTelemetry.__new__(OpenTelemetry)
otel_logger.callback_name = "ArizeLogger"
otel_logger._span_export_in_failure_state = False
with patch.object(otel_logger, "_report_callback_failure") as mock_report:
wrapped_exporter = otel_logger._wrap_exporter_with_failure_tracking(
_DummyExporter()
)
wrapped_exporter.export([])
wrapped_exporter.export([])
wrapped_exporter.export([])
wrapped_exporter.export([])
assert mock_report.call_count == 2
mock_report.assert_any_call(callback_name="ArizeLogger")
# ==============================================================================
# END CALLBACK FAILURE METRICS TESTS
# ==============================================================================