phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat - add remaining budget for key on prometheus

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-13 14:37:02 -07:00
parent 51157c148f
commit 9b340fb2f8
2 changed files with 29 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

27
litellm/integrations/prometheus.py
View file

@ -74,13 +74,20 @@ class PrometheusLogger:
], ],
) )
# Remaining Budget for Team, Key # Remaining Budget for Team
self.litellm_remaining_team_budget_metric = Gauge( self.litellm_remaining_team_budget_metric = Gauge(
"litellm_remaining_team_budget_metric", "litellm_remaining_team_budget_metric",
"Remaining budget for team", "Remaining budget for team",
labelnames=["team_id", "team_alias"], 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: except Exception as e:
print_verbose(f"Got exception on init prometheus client {str(e)}") print_verbose(f"Got exception on init prometheus client {str(e)}")
raise e raise e
@ -111,6 +118,9 @@ class PrometheusLogger:
"user_api_key_user_id", None "user_api_key_user_id", None
) )
user_api_key = litellm_params.get("metadata", {}).get("user_api_key", 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_team = litellm_params.get("metadata", {}).get(
"user_api_key_team_id", None "user_api_key_team_id", None
) )
@ -121,7 +131,6 @@ class PrometheusLogger:
_team_spend = litellm_params.get("metadata", {}).get( _team_spend = litellm_params.get("metadata", {}).get(
"user_api_key_team_spend", None "user_api_key_team_spend", None
) )
_team_max_budget = litellm_params.get("metadata", {}).get( _team_max_budget = litellm_params.get("metadata", {}).get(
"user_api_key_team_max_budget", None "user_api_key_team_max_budget", None
) )
@ -129,6 +138,16 @@ class PrometheusLogger:
max_budget=_team_max_budget, spend=_team_spend 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: if response_obj is not None:
tokens_used = response_obj.get("usage", {}).get("total_tokens", 0) tokens_used = response_obj.get("usage", {}).get("total_tokens", 0)
else: else:
@ -176,6 +195,10 @@ class PrometheusLogger:
user_api_team, user_api_team_alias user_api_team, user_api_team_alias
).set(_remaining_team_budget) ).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 ### ### FAILURE INCREMENT ###
if "exception" in kwargs: if "exception" in kwargs:
self.litellm_llm_api_failed_requests_metric.labels( self.litellm_llm_api_failed_requests_metric.labels(

4
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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_max_budget"] = user_api_key_dict.team_max_budget
data["metadata"]["user_api_key_team_spend"] = user_api_key_dict.team_spend 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 data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata
_headers = dict(request.headers) _headers = dict(request.headers)
_headers.pop( _headers.pop(