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"],
)
# 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],

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
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):
"""