diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f575372fc3d..743287c52e4 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1716,6 +1716,35 @@ class PrometheusLogger(CustomLogger): amount=float(response_cost), ) + @staticmethod + def _get_remaining_from_v3_rate_limit_headers( + standard_logging_payload: StandardLoggingPayload | None, + rate_limit_type: Literal["requests", "tokens"], + ) -> int | None: + """ + Read the per-(key, model) remaining value emitted by the v3 rate + limiter (``parallel_request_limiter_v3.py``), which writes + ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` into + ``standard_logging_object.hidden_params.additional_headers`` instead + of the ``litellm-key-remaining-*`` metadata keys the legacy limiter + sets. The header carries no model group; it always refers to this + request's model group, which is what the gauges are labeled with. + Values are written in-process as plain ints (never HTTP-serialized + strings), so anything else is rejected rather than coerced. + """ + if standard_logging_payload is None: + return None + hidden_params = standard_logging_payload.get("hidden_params") + if hidden_params is None: + return None + additional_headers = hidden_params.get("additional_headers") + if additional_headers is None: + return None + value = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}") + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + def _set_virtual_key_rate_limit_metrics( self, user_api_key: Optional[str], @@ -1733,11 +1762,20 @@ class PrometheusLogger(CustomLogger): model_group = get_model_group_from_litellm_kwargs(kwargs) remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" + standard_logging_payload: StandardLoggingPayload | None = kwargs.get("standard_logging_object") remaining_requests = metadata.get(remaining_requests_variable_name) + if remaining_requests is None: + remaining_requests = self._get_remaining_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, rate_limit_type="requests" + ) if remaining_requests is None: remaining_requests = sys.maxsize remaining_tokens = metadata.get(remaining_tokens_variable_name) + if remaining_tokens is None: + remaining_tokens = self._get_remaining_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, rate_limit_type="tokens" + ) if remaining_tokens is None: remaining_tokens = sys.maxsize diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index bb035c4c3ee..9c6d2e018ff 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -326,3 +326,148 @@ async def test_should_leave_rate_limit_labels_blank_for_non_rate_limit_failure() assert isinstance(enum_values, UserAPIKeyLabelValues) assert enum_values.rate_limit_category is None assert enum_values.rate_limit_type is None + + +def _logger_with_mock_virtual_key_gauges() -> PrometheusLogger: + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_remaining_api_key_requests_for_model = MagicMock() + logger.litellm_remaining_api_key_tokens_for_model = MagicMock() + logger.get_labels_for_metric = MagicMock(return_value=[]) + return logger + + +def _kwargs_with_v3_rate_limit_headers(additional_headers: dict) -> dict: + return { + "litellm_params": {"metadata": {"model_group": "gpt-4o-mini"}}, + "standard_logging_object": { + "metadata": {}, + "hidden_params": {"additional_headers": additional_headers}, + }, + } + + +def _set_virtual_key_metrics(logger: PrometheusLogger, kwargs: dict) -> None: + logger._set_virtual_key_rate_limit_metrics( + user_api_key="test-hash", + user_api_key_alias="test-alias", + kwargs=kwargs, + metadata=kwargs["litellm_params"]["metadata"], + model_id="model-123", + ) + + +def test_should_read_v3_remaining_headers_when_metadata_keys_absent(): + """ + Regression for LIT-2577: the default v3 rate limiter writes remaining + per-(key, model) values into + ``standard_logging_object.hidden_params.additional_headers`` as + ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` and never sets + the legacy ``litellm-key-remaining-*`` metadata keys, so the gauges were + pinned to ``sys.maxsize``. + """ + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-remaining-tokens": 900, + "x-ratelimit-model_per_key-limit-requests": 100, + "x-ratelimit-model_per_key-limit-tokens": 1000, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 42 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 900 + ) + + +def test_should_prefer_legacy_metadata_keys_over_v3_headers(): + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-remaining-tokens": 900, + } + ) + kwargs["litellm_params"]["metadata"].update( + { + "litellm-key-remaining-requests-gpt-4o-mini": 3, + "litellm-key-remaining-tokens-gpt-4o-mini": 200, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 3 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 200 + ) + + +def test_should_treat_zero_v3_remaining_as_zero(): + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 0, + "x-ratelimit-model_per_key-remaining-tokens": 0, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + 0 + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + 0 + ) + + +def test_should_keep_maxsize_sentinel_when_no_rate_limit_source_present(): + import sys + + logger = _logger_with_mock_virtual_key_gauges() + kwargs = { + "litellm_params": {"metadata": {"model_group": "gpt-4o-mini"}}, + "standard_logging_object": {"metadata": {}, "hidden_params": {}}, + } + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + + +@pytest.mark.parametrize("bad_value", ["not-a-number", None, True]) +def test_should_ignore_non_int_v3_header_values(bad_value): + import sys + + logger = _logger_with_mock_virtual_key_gauges() + kwargs = _kwargs_with_v3_rate_limit_headers( + { + "x-ratelimit-model_per_key-remaining-requests": bad_value, + "x-ratelimit-model_per_key-remaining-tokens": bad_value, + } + ) + + _set_virtual_key_metrics(logger, kwargs) + + logger.litellm_remaining_api_key_requests_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + ) + logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( + sys.maxsize + )