From 8d4ad47ec3a64ccd156a6b0399ee6062d5dd8864 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Tue, 4 Feb 2025 19:15:50 -0800 Subject: [PATCH] fix(prometheus.py): fix setting key budget metrics (#8234) * fix(prometheus.py): fix setting key budget metrics ensures custom metadata works with key budget metric this is a patch. root cause pr is written in a separate branch * test: fix test --- litellm/integrations/prometheus.py | 31 +++++++++++++------ .../test_prometheus_unit_tests.py | 6 ++-- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 01e4346afe2..4759ef91e63 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1587,10 +1587,17 @@ class PrometheusLogger(CustomLogger): - Max Budget - Budget Reset At """ - self.litellm_remaining_api_key_budget_metric.labels( - user_api_key_dict.token, - user_api_key_dict.key_alias or "", - ).set( + enum_values = UserAPIKeyLabelValues( + hashed_api_key=user_api_key_dict.token, + api_key_alias=user_api_key_dict.key_alias or "", + ) + _labels = prometheus_label_factory( + supported_enum_labels=PrometheusMetricLabels.get_labels( + label_name="litellm_remaining_api_key_budget_metric" + ), + enum_values=enum_values, + ) + self.litellm_remaining_api_key_budget_metric.labels(**_labels).set( self._safe_get_remaining_budget( max_budget=user_api_key_dict.max_budget, spend=user_api_key_dict.spend, @@ -1598,14 +1605,18 @@ class PrometheusLogger(CustomLogger): ) if user_api_key_dict.max_budget is not None: - self.litellm_api_key_max_budget_metric.labels( - user_api_key_dict.token, user_api_key_dict.key_alias - ).set(user_api_key_dict.max_budget) + _labels = prometheus_label_factory( + supported_enum_labels=PrometheusMetricLabels.get_labels( + label_name="litellm_api_key_max_budget_metric" + ), + enum_values=enum_values, + ) + self.litellm_api_key_max_budget_metric.labels(**_labels).set( + user_api_key_dict.max_budget + ) if user_api_key_dict.budget_reset_at is not None: - self.litellm_api_key_budget_remaining_hours_metric.labels( - user_api_key_dict.token, user_api_key_dict.key_alias - ).set( + self.litellm_api_key_budget_remaining_hours_metric.labels(**_labels).set( self._get_remaining_hours_for_budget_reset( budget_reset_at=user_api_key_dict.budget_reset_at ) diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index 1b157dd3356..7a64dd64eb5 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -308,7 +308,7 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): ) prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once_with( - "key1", "alias1" + hashed_api_key="key1", api_key_alias="alias1" ) prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.assert_called_once_with( 40 # 75 - (25 + 10) @@ -323,7 +323,7 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): ) prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once_with( - "key1", "alias1" + hashed_api_key="key1", api_key_alias="alias1" ) prometheus_logger.litellm_api_key_max_budget_metric.labels().set.assert_called_once_with( 75 @@ -342,7 +342,7 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): assert 9.9 <= remaining_hours_call <= 10.0 prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels.assert_called_once_with( - "key1", "alias1" + hashed_api_key="key1", api_key_alias="alias1" ) # The remaining hours should be approximately 10 (with some small difference due to test execution time) remaining_hours_call = prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.call_args[