From e6517090728ce073715947bd3f457a445d377b5a Mon Sep 17 00:00:00 2001
From: reitowo <29846655+reitowo@users.noreply.github.com>
Date: Mon, 3 Aug 2026 15:36:14 +0800
Subject: [PATCH] fix(prometheus): honor cooldown label filters

---
 litellm/integrations/prometheus.py            | 18 ++++++-----
 .../test_prometheus_logging_callbacks.py      | 31 ++++++-------------
 2 files changed, 21 insertions(+), 28 deletions(-)

diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py
index c84a6c34f1f..7756729fa2b 100644
--- a/litellm/integrations/prometheus.py
+++ b/litellm/integrations/prometheus.py
@@ -3239,13 +3239,17 @@ class PrometheusLogger(CustomLogger):
         """
         increment metric when litellm.Router / load balancing logic places a deployment in cool down
         """
-        self.litellm_deployment_cooled_down.labels(
-            _sanitize_prometheus_label_value(litellm_model_name),
-            _sanitize_prometheus_label_value(model_id),
-            _sanitize_prometheus_label_value(api_base),
-            _sanitize_prometheus_label_value(api_provider),
-            _sanitize_prometheus_label_value(exception_status),
-        ).inc()
+        self._inc_labeled_counter(
+            counter=self.litellm_deployment_cooled_down,
+            metric_name="litellm_deployment_cooled_down",
+            enum_values=UserAPIKeyLabelValues(
+                litellm_model_name=litellm_model_name,
+                model_id=model_id,
+                api_base=api_base,
+                api_provider=api_provider,
+                exception_status=exception_status,
+            ),
+        )
 
     def increment_callback_logging_failure(
         self,
diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
index 9acb87750e9..440e96e9f38 100644
--- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
+++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
@@ -1177,39 +1177,28 @@ def test_deployment_state_management(prometheus_logger):
 
 
 def test_increment_deployment_cooled_down(prometheus_logger):
-    import inspect
-
-    method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down)
-    expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"])
-
     mock_chain = MagicMock()
-
-    def validating_labels(*label_values, **label_kwargs):
-        """Validate label count matches metric definition"""
-        total = len(label_values) + len(label_kwargs)
-        if total != expected_label_count:
-            raise ValueError(
-                f"Incorrect label count: expected {expected_label_count}, got {total}"
-            )
-        return mock_chain
-
     prometheus_logger.litellm_deployment_cooled_down = MagicMock()
-    prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(
-        side_effect=validating_labels
+    prometheus_logger.litellm_deployment_cooled_down.labels.return_value = mock_chain
+    prometheus_logger.get_labels_for_metric = MagicMock(
+        return_value=["litellm_model_name", "model_id", "api_base", "api_provider"]
     )
 
     prometheus_logger.increment_deployment_cooled_down(
-        litellm_model_name="gpt-5-mini",
+        litellm_model_name="test-model",
         model_id="model-123",
-        api_base="https://api.openai.com",
+        api_base="https://api.example.com",
         api_provider="openai",
         exception_status="429",
     )
 
     prometheus_logger.litellm_deployment_cooled_down.labels.assert_called_once_with(
-        "gpt-5-mini", "model-123", "https://api.openai.com", "openai", "429"
+        litellm_model_name="test-model",
+        model_id="model-123",
+        api_base="https://api.example.com",
+        api_provider="openai",
     )
-    mock_chain.inc.assert_called_once()
+    mock_chain.inc.assert_called_once_with(1.0)
 
 
 @pytest.mark.parametrize("enable_end_user_cost_tracking_prometheus_only", [True, False])