tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 18:21:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix Prometheus custom metadata label counts (#27268) (#27271)

Browse Source 
* Fix Prometheus custom metadata label counts (#27268)

Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com>
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

* fix enterprise test: update positional label assertions to keyword args

prometheus_label_factory now calls .labels() with keyword arguments.
Update test_async_log_failure_event assertion to match.

---------

Co-authored-by: oss-agent-shin <ext-agent-shin@berri.ai>
Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com>
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
ishaan-berri
2026-05-05 20:04:56 -07:00
committed by GitHub
co-authored by oss-agent-shin oss-agent-shin ishaan-berri
parent e9fb29061a
commit c32ad90823
3 changed files with 217 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/prometheus.py
+84 -37
View File
@@ -1047,21 +1047,9 @@ class PrometheusLogger(CustomLogger):
output_tokens = standard_logging_payload["completion_tokens"]
tokens_used = standard_logging_payload["total_tokens"]
response_cost = standard_logging_payload["response_cost"]
_requester_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"requester_metadata"
combined_metadata = _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
)
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
"metadata"
].get("user_api_key_auth_metadata")
spend_logs_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"spend_logs_metadata"
)
combined_metadata: Dict[str, Any] = {
**(_requester_metadata if _requester_metadata else {}),
**(user_api_key_auth_metadata if user_api_key_auth_metadata else {}),
**(spend_logs_metadata if spend_logs_metadata else {}),
}
if standard_logging_payload is not None and isinstance(
standard_logging_payload, dict
):
@@ -1423,19 +1411,39 @@ 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)
enum_values = UserAPIKeyLabelValues(
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
model=model_group,
model_id=model_id,
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=kwargs.get("standard_logging_object")
)
),
)
label_context = PrometheusLabelFactoryContext(enum_values)
requests_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_requests_for_model"
),
enum_values=enum_values,
label_context=label_context,
)
self.litellm_remaining_api_key_requests_for_model.labels(**requests_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)
tokens_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_tokens_for_model"
),
enum_values=enum_values,
label_context=label_context,
)
self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(
remaining_tokens
)
def _set_latency_metrics(
self,
@@ -1561,18 +1569,27 @@ class PrometheusLogger(CustomLogger):
)
try:
self.litellm_llm_api_failed_requests_metric.labels(
_sanitize_prometheus_label_value(end_user_id),
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model),
_sanitize_prometheus_label_value(user_api_team),
_sanitize_prometheus_label_value(user_api_team_alias),
_sanitize_prometheus_label_value(user_id),
_sanitize_prometheus_label_value(
standard_logging_payload.get("model_id", "")
enum_values = UserAPIKeyLabelValues(
end_user=end_user_id,
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
model=model,
team=user_api_team,
team_alias=user_api_team_alias,
user=user_id,
model_id=standard_logging_payload.get("model_id", ""),
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
)
),
).inc()
)
PrometheusLogger._inc_labeled_counter(
self,
self.litellm_llm_api_failed_requests_metric,
"litellm_llm_api_failed_requests_metric",
enum_values,
)
self.set_llm_deployment_failure_metrics(kwargs)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
@@ -3622,6 +3639,36 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]:
return result
def _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload: Optional[dict],
) -> Dict[str, Any]:
"""
Combine the metadata sources that can supply custom Prometheus labels.
"""
if not isinstance(standard_logging_payload, dict):
return {}
standard_logging_metadata = standard_logging_payload.get("metadata") or {}
if not isinstance(standard_logging_metadata, dict):
return {}
requester_metadata = standard_logging_metadata.get("requester_metadata")
user_api_key_auth_metadata = standard_logging_metadata.get(
"user_api_key_auth_metadata"
)
spend_logs_metadata = standard_logging_metadata.get("spend_logs_metadata")
return {
**(requester_metadata if isinstance(requester_metadata, dict) else {}),
**(
user_api_key_auth_metadata
if isinstance(user_api_key_auth_metadata, dict)
else {}
),
**(spend_logs_metadata if isinstance(spend_logs_metadata, dict) else {}),
}
def _tag_matches_wildcard_configured_pattern(
tags: Sequence[str], configured_tag: str
) -> bool:
tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
+8 -8
View File
@@ -661,14 +661,14 @@ async def test_async_log_failure_event(prometheus_logger):
# litellm_llm_api_failed_requests_metric incremented
# Labels: end_user, hashed_api_key, api_key_alias, model, team, team_alias, user, model_id
prometheus_logger.litellm_llm_api_failed_requests_metric.labels.assert_called_once_with(
None, # end_user_id
"test_hash",
"test_alias",
"gpt-3.5-turbo",
"test_team",
"test_team_alias",
"test_user",
"model-123", # model_id from standard_logging_payload
end_user=None,
hashed_api_key="test_hash",
api_key_alias="test_alias",
model="gpt-3.5-turbo",
team="test_team",
team_alias="test_team_alias",
user="test_user",
model_id="model-123",
)
prometheus_logger.litellm_llm_api_failed_requests_metric.labels().inc.assert_called_once()
tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py
+125
View File
@@ -0,0 +1,125 @@
import logging
import pytest
from prometheus_client import REGISTRY
import litellm
from litellm.integrations.prometheus import PrometheusLogger
def _clear_prometheus_registry() -> None:
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
def _create_prometheus_logger_with_custom_labels(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
litellm,
"custom_prometheus_metadata_labels",
["metadata.department", "metadata.environment"],
)
_clear_prometheus_registry()
return PrometheusLogger()
def _standard_logging_payload_with_requester_metadata() -> dict:
return {
"model_id": "model-123",
"model_group": "gpt-4o-mini",
"api_base": "https://api.openai.com",
"custom_llm_provider": "openai",
"metadata": {
"user_api_key_hash": "test-hash",
"user_api_key_alias": "test-alias",
"user_api_key_team_id": "test-team",
"user_api_key_team_alias": "test-team-alias",
"user_api_key_user_id": "test-user",
"user_api_key_user_email": "test@example.com",
"user_api_key_org_id": None,
"requester_metadata": {
"department": "engineering",
"environment": "production",
},
"user_api_key_auth_metadata": None,
"spend_logs_metadata": None,
},
"request_tags": [],
"completion_tokens": 0,
"total_tokens": 0,
"response_cost": 0,
}
def _metric_samples(metric_name: str):
return [
sample
for metric in REGISTRY.collect()
for sample in metric.samples
if sample.name == metric_name
]
@pytest.mark.asyncio
async def test_async_log_failure_event_accepts_custom_metadata_labels(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
):
prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch)
kwargs = {
"model": "gpt-4o-mini",
"litellm_params": {
"metadata": {
"user_api_key_end_user_id": "test-end-user",
}
},
"standard_logging_object": _standard_logging_payload_with_requester_metadata(),
}
with caplog.at_level(logging.ERROR):
await prometheus_logger.async_log_failure_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
assert "Incorrect label count" not in caplog.text
samples = _metric_samples("litellm_llm_api_failed_requests_metric_total")
assert any(
sample.labels.get("metadata_department") == "engineering"
and sample.labels.get("metadata_environment") == "production"
for sample in samples
)
def test_virtual_key_rate_limit_metrics_accept_custom_metadata_labels(
monkeypatch: pytest.MonkeyPatch,
):
prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch)
metadata = {
"model_group": "gpt-4o-mini",
"litellm-key-remaining-requests-gpt-4o-mini": 3,
"litellm-key-remaining-tokens-gpt-4o-mini": 200,
}
kwargs = {
"litellm_params": {
"metadata": metadata,
},
"standard_logging_object": _standard_logging_payload_with_requester_metadata(),
}
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",
)
samples = _metric_samples("litellm_remaining_api_key_requests_for_model")
assert any(
sample.labels.get("metadata_department") == "engineering"
and sample.labels.get("metadata_environment") == "production"
and sample.value == 3
for sample in samples
)