diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 7d3856ea63..147420cd53 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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: diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 9f879b9d50..e6c76fa7cd 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -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() 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 new file mode 100644 index 0000000000..05d1eab8cc --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py @@ -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 + )