From ce9af1e1dfcc11c763b2d20a109e7188cb263bf8 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 10 Aug 2024 09:28:46 -0700
Subject: [PATCH] refactor prometheus to be a customLogger class

---
 litellm/__init__.py                           |  1 +
 litellm/litellm_core_utils/litellm_logging.py | 15 ++++++++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index 7e23e92a4a..fdfe3bd9ea 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 cef7973128..9f84b26d66 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):