diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f9b1c666439..d519d68d346 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1444,9 +1444,23 @@ class PrometheusLogger(CustomLogger): remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" remaining_requests = metadata.get(remaining_requests_variable_name) + if remaining_requests is None: + remaining_requests = ( + self._get_model_per_key_remaining_value_from_additional_headers( + kwargs=kwargs, + 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_model_per_key_remaining_value_from_additional_headers( + kwargs=kwargs, + rate_limit_type="tokens", + ) + ) if remaining_tokens is None: remaining_tokens = sys.maxsize @@ -1484,6 +1498,28 @@ class PrometheusLogger(CustomLogger): remaining_tokens ) + @staticmethod + def _get_model_per_key_remaining_value_from_additional_headers( + kwargs: dict, + rate_limit_type: Literal["requests", "tokens"], + ) -> Optional[Any]: + standard_logging_payload = kwargs.get("standard_logging_object") or {} + hidden_params = standard_logging_payload.get("hidden_params") or {} + additional_headers = hidden_params.get("additional_headers") or {} + if not isinstance(additional_headers, dict): + return None + + header_name = f"x-ratelimit-model_per_key-remaining-{rate_limit_type}" + if header_name in additional_headers: + return additional_headers[header_name] + + normalized_header_name = header_name.lower() + for key, value in additional_headers.items(): + if str(key).lower() == normalized_header_name: + return value + + return None + def _set_latency_metrics( self, kwargs: dict, diff --git a/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py b/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py index 99eb5abb7b5..8dcac51d506 100644 --- a/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py +++ b/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py @@ -157,3 +157,81 @@ def test_virtual_key_rate_limit_metrics_preserve_zero_remaining_values( assert any(sample.value == 0 for sample in token_samples) assert not any(sample.value == sys.maxsize for sample in request_samples) assert not any(sample.value == sys.maxsize for sample in token_samples) + + +def test_virtual_key_rate_limit_metrics_read_remaining_values_from_hidden_headers( + monkeypatch: pytest.MonkeyPatch, +): + prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch) + metadata = { + "model_group": "gpt-4o-mini", + } + kwargs = { + "litellm_params": { + "metadata": metadata, + }, + "standard_logging_object": { + **_standard_logging_payload_with_requester_metadata(), + "hidden_params": { + "additional_headers": { + "x-ratelimit-model_per_key-remaining-requests": 7, + "x-ratelimit-model_per_key-remaining-tokens": 321, + }, + }, + }, + } + + prometheus_logger._set_virtual_key_rate_limit_metrics( + user_api_key="test-hash", + user_api_key_alias="test-alias", + kwargs=kwargs, + metadata=metadata, + model_id="model-123", + ) + + request_samples = _metric_samples("litellm_remaining_api_key_requests_for_model") + token_samples = _metric_samples("litellm_remaining_api_key_tokens_for_model") + + assert any(sample.value == 7 for sample in request_samples) + assert any(sample.value == 321 for sample in token_samples) + assert not any(sample.value == sys.maxsize for sample in request_samples) + assert not any(sample.value == sys.maxsize for sample in token_samples) + + +def test_virtual_key_rate_limit_metrics_preserve_zero_values_from_hidden_headers( + monkeypatch: pytest.MonkeyPatch, +): + prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch) + metadata = { + "model_group": "gpt-4o-mini", + } + kwargs = { + "litellm_params": { + "metadata": metadata, + }, + "standard_logging_object": { + **_standard_logging_payload_with_requester_metadata(), + "hidden_params": { + "additional_headers": { + "x-ratelimit-model_per_key-remaining-requests": 0, + "x-ratelimit-model_per_key-remaining-tokens": 0, + }, + }, + }, + } + + prometheus_logger._set_virtual_key_rate_limit_metrics( + user_api_key="test-hash", + user_api_key_alias="test-alias", + kwargs=kwargs, + metadata=metadata, + model_id="model-123", + ) + + request_samples = _metric_samples("litellm_remaining_api_key_requests_for_model") + token_samples = _metric_samples("litellm_remaining_api_key_tokens_for_model") + + assert any(sample.value == 0 for sample in request_samples) + assert any(sample.value == 0 for sample in token_samples) + assert not any(sample.value == sys.maxsize for sample in request_samples) + assert not any(sample.value == sys.maxsize for sample in token_samples)