scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

Merge c560859f39 into 26fcbc93e5

Browse source
This commit is contained in:
tom-onyx 2026-04-19 09:25:12 +03:00 • committed by GitHub
commit c06983c2be
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 203 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/integrations/prometheus.py
View file

@ -1165,11 +1165,9 @@ class PrometheusLogger(CustomLogger):
# set proxy virtual key rpm/tpm metrics
self._set_virtual_key_rate_limit_metrics(
user_api_key=user_api_key,
user_api_key_alias=user_api_key_alias,
kwargs=kwargs,
metadata=_metadata,
model_id=enum_values.model_id,
enum_values=enum_values,
)
# set latency metrics
@ -1396,11 +1394,9 @@ class PrometheusLogger(CustomLogger):
def _set_virtual_key_rate_limit_metrics(
self,
user_api_key: Optional[str],
user_api_key_alias: Optional[str],
kwargs: dict,
metadata: dict,
model_id: Optional[str] = None,
enum_values: UserAPIKeyLabelValues,
):
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
@ -1421,19 +1417,25 @@ class PrometheusLogger(CustomLogger):
metadata.get(remaining_tokens_variable_name, sys.maxsize) or sys.maxsize
)
self.litellm_remaining_api_key_requests_for_model.labels(
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model_group),
_sanitize_prometheus_label_value(model_id),
).set(remaining_requests)
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_requests_for_model"
),
enum_values=enum_values,
)
self.litellm_remaining_api_key_requests_for_model.labels(**_labels).set(
remaining_requests
)
self.litellm_remaining_api_key_tokens_for_model.labels(
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model_group),
_sanitize_prometheus_label_value(model_id),
).set(remaining_tokens)
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_tokens_for_model"
),
enum_values=enum_values,
)
self.litellm_remaining_api_key_tokens_for_model.labels(**_labels).set(
remaining_tokens
)
def _set_latency_metrics(
self,

183
tests/test_litellm/integrations/test_prometheus_virtual_key_rate_limit.py Normal file
View file

@ -0,0 +1,183 @@
"""
Unit tests for _set_virtual_key_rate_limit_metrics using prometheus_label_factory.
When custom_prometheus_metadata_labels is configured (e.g., ["onyx_feature"]),
PrometheusMetricLabels.get_labels() appends the custom labels to every metric's
label set. The Gauge is created with these extra labels, so callers must provide
values for them. Using hardcoded positional args to .labels() fails with
ValueError: Incorrect label count
because the custom labels are never passed.
The fix uses prometheus_label_factory (same pattern as all other metrics in the
callback) so that custom metadata labels are resolved from enum_values
automatically.
Fixes https://github.com/BerriAI/litellm/issues/24760
"""
import os
import sys
import pytest
from prometheus_client import REGISTRY
sys.path.insert(0, os.path.abspath("../../.."))
import litellm
from litellm.integrations.prometheus import PrometheusLogger
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
@pytest.fixture(scope="function")
def prometheus_logger():
"""Create a PrometheusLogger instance for testing."""
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
return PrometheusLogger()
@pytest.fixture(scope="function")
def prometheus_logger_with_custom_labels():
"""Create a PrometheusLogger with custom_prometheus_metadata_labels configured."""
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
original = litellm.custom_prometheus_metadata_labels
litellm.custom_prometheus_metadata_labels = ["onyx_feature"]
try:
logger = PrometheusLogger()
yield logger
finally:
litellm.custom_prometheus_metadata_labels = original
class TestVirtualKeyRateLimitMetrics:
"""
Test that _set_virtual_key_rate_limit_metrics works correctly with
prometheus_label_factory, including when custom_prometheus_metadata_labels
is configured.
"""
def test_set_virtual_key_rate_limit_metrics_basic(self, prometheus_logger):
"""
_set_virtual_key_rate_limit_metrics should not raise with default labels.
"""
enum_values = UserAPIKeyLabelValues(
hashed_api_key="test-key-hash",
api_key_alias="test-alias",
model="gpt-4o",
model_id="model-123",
)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-4o",
},
},
}
metadata = {}
# Should not raise
prometheus_logger._set_virtual_key_rate_limit_metrics(
kwargs=kwargs,
metadata=metadata,
enum_values=enum_values,
)
def test_set_virtual_key_rate_limit_metrics_with_custom_labels(
self, prometheus_logger_with_custom_labels
):
"""
_set_virtual_key_rate_limit_metrics should not raise ValueError when
custom_prometheus_metadata_labels adds extra labels to the Gauge.
Before the fix, this crashed with:
ValueError: Incorrect label count
because the Gauge had 5 labels (4 default + 1 custom) but only 4
positional args were passed.
"""
enum_values = UserAPIKeyLabelValues(
hashed_api_key="test-key-hash",
api_key_alias="test-alias",
model="gpt-4o",
model_id="model-123",
custom_metadata_labels={"onyx_feature": "test-feature"},
)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-4o",
},
},
}
metadata = {}
# Should not raise ValueError: Incorrect label count
prometheus_logger_with_custom_labels._set_virtual_key_rate_limit_metrics(
kwargs=kwargs,
metadata=metadata,
enum_values=enum_values,
)
def test_set_virtual_key_rate_limit_metrics_with_none_values(
self, prometheus_logger
):
"""
_set_virtual_key_rate_limit_metrics should handle None values gracefully.
"""
enum_values = UserAPIKeyLabelValues(
hashed_api_key=None,
api_key_alias=None,
model=None,
model_id=None,
)
kwargs = {
"litellm_params": {
"metadata": {},
},
}
metadata = {}
# Should not raise
prometheus_logger._set_virtual_key_rate_limit_metrics(
kwargs=kwargs,
metadata=metadata,
enum_values=enum_values,
)
def test_set_virtual_key_rate_limit_metrics_sets_remaining_values(
self, prometheus_logger
):
"""
_set_virtual_key_rate_limit_metrics should correctly set the remaining
request/token counts from metadata.
"""
enum_values = UserAPIKeyLabelValues(
hashed_api_key="test-key-hash",
api_key_alias="test-alias",
model="gpt-4o",
model_id="model-123",
)
kwargs = {
"litellm_params": {
"metadata": {
"model_group": "gpt-4o",
},
},
}
metadata = {
"litellm-key-remaining-requests-gpt-4o": 100,
"litellm-key-remaining-tokens-gpt-4o": 50000,
}
# Should not raise
prometheus_logger._set_virtual_key_rate_limit_metrics(
kwargs=kwargs,
metadata=metadata,
enum_values=enum_values,
)