ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(prometheus.py): update tests

Browse Source
This commit is contained in:
Krrish Dholakia 2025-06-06 09:12:54 -07:00
parent 1b7056f281
commit 3608db5ffe
2 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
litellm/integrations/prometheus.py
View File

@ -543,6 +543,7 @@ class PrometheusLogger(CustomLogger):
user_id: Optional[str],
enum_values: UserAPIKeyLabelValues,
):
verbose_logger.debug(f"prometheus Logging - Enters token metrics function")
# token metrics
if standard_logging_payload is not None and isinstance(
@ -556,10 +557,21 @@ class PrometheusLogger(CustomLogger):
),
enum_values=enum_values,
)
self.litellm_proxy_total_requests_metric.labels(**_labels).inc(
standard_logging_payload["total_tokens"]
)
_labels = prometheus_label_factory(
supported_enum_labels=PrometheusMetricLabels.get_labels(
label_name="litellm_total_tokens_metric"
),
enum_values=enum_values,
)
self.litellm_tokens_metric.labels(**_labels).inc(
standard_logging_payload["total_tokens"]
)
_labels = prometheus_label_factory(
supported_enum_labels=PrometheusMetricLabels.get_labels(
label_name="litellm_input_tokens_metric"

2
tests/logging_callback_tests/test_prometheus_unit_tests.py
View File

@ -219,7 +219,7 @@ def test_increment_token_metrics(prometheus_logger):
)
prometheus_logger.litellm_tokens_metric.labels.assert_called_once_with(
"user1", "key1", "alias1", "gpt-3.5-turbo", "team1", "team_alias1", "user1"
end_user=None, user=None, hashed_api_key='test_hash', api_key_alias='test_alias', team='test_team', team_alias='test_team_alias', requested_model=None, model='gpt-3.5-turbo'
)
prometheus_logger.litellm_tokens_metric.labels().inc.assert_called_once_with(100)