phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add key and team level budget (#7831)

Browse source
This commit is contained in:
yujonglee 2025-01-18 02:04:12 +09:00 • committed by GitHub
parent 4538df580e
commit b0e30906e0
2 changed files with 44 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

26
litellm/integrations/prometheus.py
View file

@ -149,6 +149,20 @@ class PrometheusLogger(CustomLogger):
labelnames=["hashed_api_key", "api_key_alias"], 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 # LiteLLM Virtual API KEY metrics
######################################## ########################################
@ -557,10 +571,22 @@ class PrometheusLogger(CustomLogger):
user_api_team, user_api_team_alias user_api_team, user_api_team_alias
).set(_remaining_team_budget) ).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( self.litellm_remaining_api_key_budget_metric.labels(
user_api_key, user_api_key_alias user_api_key, user_api_key_alias
).set(_remaining_api_key_budget) ).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( def _increment_top_level_request_and_spend_metrics(
self, self,
end_user_id: Optional[str], end_user_id: Optional[str],

21
tests/logging_callback_tests/test_prometheus_unit_tests.py
View file

@ -255,10 +255,13 @@ def test_increment_remaining_budget_metrics(prometheus_logger):
""" """
Test the increment_remaining_budget_metrics method 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_team_budget_metric = MagicMock()
prometheus_logger.litellm_remaining_api_key_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 = { litellm_params = {
"metadata": { "metadata": {
@ -277,20 +280,32 @@ def test_increment_remaining_budget_metrics(prometheus_logger):
litellm_params=litellm_params, litellm_params=litellm_params,
) )
# Test remaining budget metrics
prometheus_logger.litellm_remaining_team_budget_metric.labels.assert_called_once_with( prometheus_logger.litellm_remaining_team_budget_metric.labels.assert_called_once_with(
"team1", "team_alias1" "team1", "team_alias1"
) )
prometheus_logger.litellm_remaining_team_budget_metric.labels().set.assert_called_once_with( 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( prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once_with(
"key1", "alias1" "key1", "alias1"
) )
prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.assert_called_once_with( 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): def test_set_latency_metrics(prometheus_logger):
""" """