Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions litellm/integrations/prometheus.py
Original file line number Diff line number Diff line change
Expand Up @@ -1716,6 +1716,35 @@ def _increment_top_level_request_and_spend_metrics(
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],
Expand All @@ -1733,11 +1762,20 @@ def _set_virtual_key_rate_limit_metrics(
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

Expand Down
145 changes: 145 additions & 0 deletions tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Loading