diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index d3283cc74e6..15583bb62b7 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -324,7 +324,9 @@ class PrometheusLogger(CustomLogger): self.litellm_provider_remaining_budget_metric = self._gauge_factory( "litellm_provider_remaining_budget_metric", "Remaining budget for provider - used when you set provider budget limits", - labelnames=["api_provider"], + labelnames=self.get_labels_for_metric( + "litellm_provider_remaining_budget_metric" + ), ) # Metric for deployment state @@ -2459,7 +2461,12 @@ class PrometheusLogger(CustomLogger): """ Track provider remaining budget in Prometheus """ - self.litellm_provider_remaining_budget_metric.labels(provider).set( + self.litellm_provider_remaining_budget_metric.labels( + **self._safe_labels( + self.litellm_provider_remaining_budget_metric, + api_provider=provider, + ) + ).set( self._safe_get_remaining_budget( max_budget=budget_limit, spend=spend, diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index e04bef6ca18..ad2d9f1d7fd 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -238,6 +238,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", "litellm_in_flight_requests", + "litellm_provider_remaining_budget_metric", ] @@ -530,6 +531,10 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, ] + litellm_provider_remaining_budget_metric = [ + UserAPIKeyLabelNames.API_PROVIDER.value, + ] + litellm_remaining_api_key_requests_for_model = [ UserAPIKeyLabelNames.API_KEY_HASH.value, UserAPIKeyLabelNames.API_KEY_ALIAS.value,