diff --git a/litellm/__init__.py b/litellm/__init__.py index 7e23e92a4..fdfe3bd9e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -43,6 +43,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "logfire", "dynamic_rate_limiter", "langsmith", + "prometheus", "galileo", "braintrust", "arize", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cef797312..9f84b26d6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1904,9 +1904,6 @@ def set_callbacks(callback_list, function_id=None): openMeterLogger = OpenMeterLogger() elif callback == "datadog": dataDogLogger = DataDogLogger() - elif callback == "prometheus": - if prometheusLogger is None: - prometheusLogger = PrometheusLogger() elif callback == "dynamodb": dynamoLogger = DyanmoDBLogger() elif callback == "s3": @@ -1980,6 +1977,14 @@ def _init_custom_logger_compatible_class( _langsmith_logger = LangsmithLogger() _in_memory_loggers.append(_langsmith_logger) return _langsmith_logger # type: ignore + elif logging_integration == "prometheus": + for callback in _in_memory_loggers: + if isinstance(callback, PrometheusLogger): + return callback # type: ignore + + _prometheus_logger = PrometheusLogger() + _in_memory_loggers.append(_prometheus_logger) + return _prometheus_logger # type: ignore elif logging_integration == "gcs_bucket": for callback in _in_memory_loggers: if isinstance(callback, GCSBucketLogger): @@ -2102,6 +2107,10 @@ def get_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, LangsmithLogger): return callback + elif logging_integration == "prometheus": + for callback in _in_memory_loggers: + if isinstance(callback, PrometheusLogger): + return callback elif logging_integration == "gcs_bucket": for callback in _in_memory_loggers: if isinstance(callback, GCSBucketLogger):