From 3608db5ffe6b64d18ec799e9be395160c40d42fc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 6 Jun 2025 09:12:54 -0700 Subject: [PATCH] fix(prometheus.py): update tests --- litellm/integrations/prometheus.py | 12 ++++++++++++ .../test_prometheus_unit_tests.py | 2 +- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index c33dcc5b553..103d4de7c2d 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -543,6 +543,7 @@ class PrometheusLogger(CustomLogger): user_id: Optional[str], enum_values: UserAPIKeyLabelValues, ): + verbose_logger.debug(f"prometheus Logging - Enters token metrics function") # token metrics if standard_logging_payload is not None and isinstance( @@ -556,10 +557,21 @@ class PrometheusLogger(CustomLogger): ), enum_values=enum_values, ) + self.litellm_proxy_total_requests_metric.labels(**_labels).inc( standard_logging_payload["total_tokens"] ) + _labels = prometheus_label_factory( + supported_enum_labels=PrometheusMetricLabels.get_labels( + label_name="litellm_total_tokens_metric" + ), + enum_values=enum_values, + ) + self.litellm_tokens_metric.labels(**_labels).inc( + standard_logging_payload["total_tokens"] + ) + _labels = prometheus_label_factory( supported_enum_labels=PrometheusMetricLabels.get_labels( label_name="litellm_input_tokens_metric" diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index ade57cc1017..5b40166286f 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -219,7 +219,7 @@ def test_increment_token_metrics(prometheus_logger): ) prometheus_logger.litellm_tokens_metric.labels.assert_called_once_with( - "user1", "key1", "alias1", "gpt-3.5-turbo", "team1", "team_alias1", "user1" + end_user=None, user=None, hashed_api_key='test_hash', api_key_alias='test_alias', team='test_team', team_alias='test_team_alias', requested_model=None, model='gpt-3.5-turbo' ) prometheus_logger.litellm_tokens_metric.labels().inc.assert_called_once_with(100)