scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

Fix Prometheus remaining rate limit metrics

Browse source 
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin 2026-05-07 00:51:03 +00:00
parent a67b7a7e87
commit 416de3e275
No known key found for this signature in database
2 changed files with 114 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
litellm/integrations/prometheus.py
View file

@ -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,

78
tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py
View file

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