diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index ce57309a0c..0e35c4e840 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -149,6 +149,20 @@ class PrometheusLogger(CustomLogger): labelnames=["hashed_api_key", "api_key_alias"], ) + # Max Budget for Team + self.litellm_team_max_budget_metric = Gauge( + "litellm_team_max_budget_metric", + "Maximum budget set for team", + labelnames=["team_id", "team_alias"], + ) + + # Max Budget for API Key + self.litellm_api_key_max_budget_metric = Gauge( + "litellm_api_key_max_budget_metric", + "Maximum budget set for api key", + labelnames=["hashed_api_key", "api_key_alias"], + ) + ######################################## # LiteLLM Virtual API KEY metrics ######################################## @@ -557,10 +571,22 @@ class PrometheusLogger(CustomLogger): user_api_team, user_api_team_alias ).set(_remaining_team_budget) + # Max Budget Metrics + if _team_max_budget is not None: + self.litellm_team_max_budget_metric.labels( + user_api_team, user_api_team_alias + ).set(_team_max_budget) + self.litellm_remaining_api_key_budget_metric.labels( user_api_key, user_api_key_alias ).set(_remaining_api_key_budget) + # Max Budget Metrics for API Key + if _api_key_max_budget is not None: + self.litellm_api_key_max_budget_metric.labels( + user_api_key, user_api_key_alias + ).set(_api_key_max_budget) + def _increment_top_level_request_and_spend_metrics( self, end_user_id: Optional[str], diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index 1026368dc6..ac8a2d33e2 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -255,10 +255,13 @@ def test_increment_remaining_budget_metrics(prometheus_logger): """ Test the increment_remaining_budget_metrics method - team and api key budget metrics are set to the difference between max budget and spend + - team and api key remaining budget metrics are set to the difference between max budget and spend + - team and api key max budget metrics are set to their respective max budgets """ prometheus_logger.litellm_remaining_team_budget_metric = MagicMock() prometheus_logger.litellm_remaining_api_key_budget_metric = MagicMock() + prometheus_logger.litellm_team_max_budget_metric = MagicMock() + prometheus_logger.litellm_api_key_max_budget_metric = MagicMock() litellm_params = { "metadata": { @@ -277,20 +280,32 @@ def test_increment_remaining_budget_metrics(prometheus_logger): litellm_params=litellm_params, ) + # Test remaining budget metrics prometheus_logger.litellm_remaining_team_budget_metric.labels.assert_called_once_with( "team1", "team_alias1" ) prometheus_logger.litellm_remaining_team_budget_metric.labels().set.assert_called_once_with( - 50 + 50 # 100 - 50 ) prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once_with( "key1", "alias1" ) prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.assert_called_once_with( - 50 + 50 # 75 - 25 ) + # Test max budget metrics + prometheus_logger.litellm_team_max_budget_metric.labels.assert_called_once_with( + "team1", "team_alias1" + ) + prometheus_logger.litellm_team_max_budget_metric.labels().set.assert_called_once_with(100) + + prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once_with( + "key1", "alias1" + ) + prometheus_logger.litellm_api_key_max_budget_metric.labels().set.assert_called_once_with(75) + def test_set_latency_metrics(prometheus_logger): """