scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(prometheus.py): fix setting key budget metrics (#8234)

Browse source 
* fix(prometheus.py): fix setting key budget metrics

ensures custom metadata works with key budget metric

this is a patch. root cause pr is written in a separate branch

* test: fix test
This commit is contained in:
Krish Dholakia 2025-02-04 19:15:50 -08:00 • committed by GitHub
parent c494b2cb85
commit 8d4ad47ec3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 24 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/integrations/prometheus.py
View file

@ -1587,10 +1587,17 @@ class PrometheusLogger(CustomLogger):
- Max Budget
- Budget Reset At
"""
self.litellm_remaining_api_key_budget_metric.labels(
user_api_key_dict.token,
user_api_key_dict.key_alias or "",
).set(
enum_values = UserAPIKeyLabelValues(
hashed_api_key=user_api_key_dict.token,
api_key_alias=user_api_key_dict.key_alias or "",
)
_labels = prometheus_label_factory(
supported_enum_labels=PrometheusMetricLabels.get_labels(
label_name="litellm_remaining_api_key_budget_metric"
),
enum_values=enum_values,
)
self.litellm_remaining_api_key_budget_metric.labels(**_labels).set(
self._safe_get_remaining_budget(
max_budget=user_api_key_dict.max_budget,
spend=user_api_key_dict.spend,
@ -1598,14 +1605,18 @@ class PrometheusLogger(CustomLogger):
)
if user_api_key_dict.max_budget is not None:
self.litellm_api_key_max_budget_metric.labels(
user_api_key_dict.token, user_api_key_dict.key_alias
).set(user_api_key_dict.max_budget)
_labels = prometheus_label_factory(
supported_enum_labels=PrometheusMetricLabels.get_labels(
label_name="litellm_api_key_max_budget_metric"
),
enum_values=enum_values,
)
self.litellm_api_key_max_budget_metric.labels(**_labels).set(
user_api_key_dict.max_budget
)
if user_api_key_dict.budget_reset_at is not None:
self.litellm_api_key_budget_remaining_hours_metric.labels(
user_api_key_dict.token, user_api_key_dict.key_alias
).set(
self.litellm_api_key_budget_remaining_hours_metric.labels(**_labels).set(
self._get_remaining_hours_for_budget_reset(
budget_reset_at=user_api_key_dict.budget_reset_at
)

6
tests/logging_callback_tests/test_prometheus_unit_tests.py
View file

@ -308,7 +308,7 @@ async def test_increment_remaining_budget_metrics(prometheus_logger):
)
prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once_with(
"key1", "alias1"
hashed_api_key="key1", api_key_alias="alias1"
)
prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.assert_called_once_with(
40 # 75 - (25 + 10)
@ -323,7 +323,7 @@ async def test_increment_remaining_budget_metrics(prometheus_logger):
)
prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once_with(
"key1", "alias1"
hashed_api_key="key1", api_key_alias="alias1"
)
prometheus_logger.litellm_api_key_max_budget_metric.labels().set.assert_called_once_with(
75
@ -342,7 +342,7 @@ async def test_increment_remaining_budget_metrics(prometheus_logger):
assert 9.9 <= remaining_hours_call <= 10.0
prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels.assert_called_once_with(
"key1", "alias1"
hashed_api_key="key1", api_key_alias="alias1"
)
# The remaining hours should be approximately 10 (with some small difference due to test execution time)
remaining_hours_call = prometheus_logger.litellm_api_key_budget_remaining_hours_metric.labels().set.call_args[