fix(prometheus): read v3 rate limiter remaining values for per-key model gauges (#33119)

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
 

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
+    )