diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 7de8088c6..4f0ffa387 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -74,13 +74,20 @@ class PrometheusLogger: ], ) - # Remaining Budget for Team, Key + # Remaining Budget for Team self.litellm_remaining_team_budget_metric = Gauge( "litellm_remaining_team_budget_metric", "Remaining budget for team", labelnames=["team_id", "team_alias"], ) + # Remaining Budget for API Key + self.litellm_remaining_api_key_budget_metric = Gauge( + "litellm_remaining_api_key_budget_metric", + "Remaining budget for api key", + labelnames=["hashed_api_key", "api_key_alias"], + ) + except Exception as e: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e @@ -111,6 +118,9 @@ class PrometheusLogger: "user_api_key_user_id", None ) user_api_key = litellm_params.get("metadata", {}).get("user_api_key", None) + user_api_key_alias = litellm_params.get("metadata", {}).get( + "user_api_key_alias", None + ) user_api_team = litellm_params.get("metadata", {}).get( "user_api_key_team_id", None ) @@ -121,7 +131,6 @@ class PrometheusLogger: _team_spend = litellm_params.get("metadata", {}).get( "user_api_key_team_spend", None ) - _team_max_budget = litellm_params.get("metadata", {}).get( "user_api_key_team_max_budget", None ) @@ -129,6 +138,16 @@ class PrometheusLogger: max_budget=_team_max_budget, spend=_team_spend ) + _api_key_spend = litellm_params.get("metadata", {}).get( + "user_api_key_spend", None + ) + _api_key_max_budget = litellm_params.get("metadata", {}).get( + "user_api_key_max_budget", None + ) + _remaining_api_key_budget = safe_get_remaining_budget( + max_budget=_api_key_max_budget, spend=_api_key_spend + ) + if response_obj is not None: tokens_used = response_obj.get("usage", {}).get("total_tokens", 0) else: @@ -176,6 +195,10 @@ class PrometheusLogger: user_api_team, user_api_team_alias ).set(_remaining_team_budget) + self.litellm_remaining_api_key_budget_metric.labels( + user_api_key, user_api_key_alias + ).set(_remaining_api_key_budget) + ### FAILURE INCREMENT ### if "exception" in kwargs: self.litellm_llm_api_failed_requests_metric.labels( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fc6bd7044..5aca61a96 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -109,6 +109,10 @@ async def add_litellm_data_to_request( data["metadata"]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget data["metadata"]["user_api_key_team_spend"] = user_api_key_dict.team_spend + # API Key spend, budget - used by prometheus.py + data["metadata"]["user_api_key_spend"] = user_api_key_dict.spend + data["metadata"]["user_api_key_max_budget"] = user_api_key_dict.max_budget + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata _headers = dict(request.headers) _headers.pop(