From c560859f399209273ad38483c552887ca9b325d0 Mon Sep 17 00:00:00 2001 From: Tom Shlomi Date: Sun, 5 Apr 2026 19:08:09 +0300 Subject: [PATCH] fix(prometheus): use prometheus_label_factory in _set_virtual_key_rate_limit_metrics Fixes #24760 _set_virtual_key_rate_limit_metrics used hardcoded positional args to .labels(), which crashes with ValueError: Incorrect label count when custom_prometheus_metadata_labels is configured. The Gauge is created with labels from get_labels_for_metric() which includes custom labels, but the hardcoded call never provides values for them. Switch to prometheus_label_factory + enum_values, consistent with every other metric in the callback. This also passes enum_values from the caller instead of individual parameters. Co-Authored-By: Claude Opus 4.6 (1M context) --- litellm/integrations/prometheus.py | 38 ++-- .../test_prometheus_virtual_key_rate_limit.py | 183 ++++++++++++++++++ 2 files changed, 203 insertions(+), 18 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_virtual_key_rate_limit.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fb5fc253ae4..d2357c8a843 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1144,11 +1144,9 @@ class PrometheusLogger(CustomLogger): # set proxy virtual key rpm/tpm metrics self._set_virtual_key_rate_limit_metrics( - user_api_key=user_api_key, - user_api_key_alias=user_api_key_alias, kwargs=kwargs, metadata=_metadata, - model_id=enum_values.model_id, + enum_values=enum_values, ) # set latency metrics @@ -1372,11 +1370,9 @@ class PrometheusLogger(CustomLogger): def _set_virtual_key_rate_limit_metrics( self, - user_api_key: Optional[str], - user_api_key_alias: Optional[str], kwargs: dict, metadata: dict, - model_id: Optional[str] = None, + enum_values: UserAPIKeyLabelValues, ): from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, @@ -1397,19 +1393,25 @@ class PrometheusLogger(CustomLogger): metadata.get(remaining_tokens_variable_name, sys.maxsize) or sys.maxsize ) - self.litellm_remaining_api_key_requests_for_model.labels( - _sanitize_prometheus_label_value(user_api_key), - _sanitize_prometheus_label_value(user_api_key_alias), - _sanitize_prometheus_label_value(model_group), - _sanitize_prometheus_label_value(model_id), - ).set(remaining_requests) + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + "litellm_remaining_api_key_requests_for_model" + ), + enum_values=enum_values, + ) + self.litellm_remaining_api_key_requests_for_model.labels(**_labels).set( + remaining_requests + ) - self.litellm_remaining_api_key_tokens_for_model.labels( - _sanitize_prometheus_label_value(user_api_key), - _sanitize_prometheus_label_value(user_api_key_alias), - _sanitize_prometheus_label_value(model_group), - _sanitize_prometheus_label_value(model_id), - ).set(remaining_tokens) + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + "litellm_remaining_api_key_tokens_for_model" + ), + enum_values=enum_values, + ) + self.litellm_remaining_api_key_tokens_for_model.labels(**_labels).set( + remaining_tokens + ) def _set_latency_metrics( self, diff --git a/tests/test_litellm/integrations/test_prometheus_virtual_key_rate_limit.py b/tests/test_litellm/integrations/test_prometheus_virtual_key_rate_limit.py new file mode 100644 index 00000000000..25309a89b06 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_virtual_key_rate_limit.py @@ -0,0 +1,183 @@ +""" +Unit tests for _set_virtual_key_rate_limit_metrics using prometheus_label_factory. + +When custom_prometheus_metadata_labels is configured (e.g., ["onyx_feature"]), +PrometheusMetricLabels.get_labels() appends the custom labels to every metric's +label set. The Gauge is created with these extra labels, so callers must provide +values for them. Using hardcoded positional args to .labels() fails with + ValueError: Incorrect label count +because the custom labels are never passed. + +The fix uses prometheus_label_factory (same pattern as all other metrics in the +callback) so that custom metadata labels are resolved from enum_values +automatically. + +Fixes https://github.com/BerriAI/litellm/issues/24760 +""" + +import os +import sys + +import pytest +from prometheus_client import REGISTRY + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +@pytest.fixture(scope="function") +def prometheus_logger(): + """Create a PrometheusLogger instance for testing.""" + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + return PrometheusLogger() + + +@pytest.fixture(scope="function") +def prometheus_logger_with_custom_labels(): + """Create a PrometheusLogger with custom_prometheus_metadata_labels configured.""" + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + original = litellm.custom_prometheus_metadata_labels + litellm.custom_prometheus_metadata_labels = ["onyx_feature"] + try: + logger = PrometheusLogger() + yield logger + finally: + litellm.custom_prometheus_metadata_labels = original + + +class TestVirtualKeyRateLimitMetrics: + """ + Test that _set_virtual_key_rate_limit_metrics works correctly with + prometheus_label_factory, including when custom_prometheus_metadata_labels + is configured. + """ + + def test_set_virtual_key_rate_limit_metrics_basic(self, prometheus_logger): + """ + _set_virtual_key_rate_limit_metrics should not raise with default labels. + """ + enum_values = UserAPIKeyLabelValues( + hashed_api_key="test-key-hash", + api_key_alias="test-alias", + model="gpt-4o", + model_id="model-123", + ) + + kwargs = { + "litellm_params": { + "metadata": { + "model_group": "gpt-4o", + }, + }, + } + metadata = {} + + # Should not raise + prometheus_logger._set_virtual_key_rate_limit_metrics( + kwargs=kwargs, + metadata=metadata, + enum_values=enum_values, + ) + + def test_set_virtual_key_rate_limit_metrics_with_custom_labels( + self, prometheus_logger_with_custom_labels + ): + """ + _set_virtual_key_rate_limit_metrics should not raise ValueError when + custom_prometheus_metadata_labels adds extra labels to the Gauge. + + Before the fix, this crashed with: + ValueError: Incorrect label count + because the Gauge had 5 labels (4 default + 1 custom) but only 4 + positional args were passed. + """ + enum_values = UserAPIKeyLabelValues( + hashed_api_key="test-key-hash", + api_key_alias="test-alias", + model="gpt-4o", + model_id="model-123", + custom_metadata_labels={"onyx_feature": "test-feature"}, + ) + + kwargs = { + "litellm_params": { + "metadata": { + "model_group": "gpt-4o", + }, + }, + } + metadata = {} + + # Should not raise ValueError: Incorrect label count + prometheus_logger_with_custom_labels._set_virtual_key_rate_limit_metrics( + kwargs=kwargs, + metadata=metadata, + enum_values=enum_values, + ) + + def test_set_virtual_key_rate_limit_metrics_with_none_values( + self, prometheus_logger + ): + """ + _set_virtual_key_rate_limit_metrics should handle None values gracefully. + """ + enum_values = UserAPIKeyLabelValues( + hashed_api_key=None, + api_key_alias=None, + model=None, + model_id=None, + ) + + kwargs = { + "litellm_params": { + "metadata": {}, + }, + } + metadata = {} + + # Should not raise + prometheus_logger._set_virtual_key_rate_limit_metrics( + kwargs=kwargs, + metadata=metadata, + enum_values=enum_values, + ) + + def test_set_virtual_key_rate_limit_metrics_sets_remaining_values( + self, prometheus_logger + ): + """ + _set_virtual_key_rate_limit_metrics should correctly set the remaining + request/token counts from metadata. + """ + enum_values = UserAPIKeyLabelValues( + hashed_api_key="test-key-hash", + api_key_alias="test-alias", + model="gpt-4o", + model_id="model-123", + ) + + kwargs = { + "litellm_params": { + "metadata": { + "model_group": "gpt-4o", + }, + }, + } + metadata = { + "litellm-key-remaining-requests-gpt-4o": 100, + "litellm-key-remaining-tokens-gpt-4o": 50000, + } + + # Should not raise + prometheus_logger._set_virtual_key_rate_limit_metrics( + kwargs=kwargs, + metadata=metadata, + enum_values=enum_values, + )