From 04c348e7bb1a65078052fea4085ebeb586d5f5fb Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Tue, 3 Feb 2026 08:50:14 +0530 Subject: [PATCH] fixes failure metrics labels (#20152) Co-authored-by: Krish Dholakia --- litellm/integrations/prometheus.py | 169 ++++++++++++++++-- .../test_prometheus_logging_callbacks.py | 144 ++++++++------- 2 files changed, 230 insertions(+), 83 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2c897cb069..00c38eac18 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1683,6 +1683,108 @@ class PrometheusLogger(CustomLogger): ) pass + def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any: + """Get value from dict or Pydantic model.""" + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + def _extract_deployment_failure_label_values( + self, request_kwargs: dict + ) -> Dict[str, Optional[str]]: + """ + Extract label values for deployment failure metrics from all available + sources in request_kwargs. Falls back to litellm_params metadata and + user_api_key_auth when standard_logging_payload has None values. + """ + standard_logging_payload = ( + request_kwargs.get("standard_logging_object", {}) or {} + ) + _litellm_params = request_kwargs.get("litellm_params", {}) or {} + _metadata_raw = self._safe_get(standard_logging_payload, "metadata") or {} + if isinstance(_metadata_raw, dict): + _metadata = _metadata_raw + else: + _metadata = { + "user_api_key_alias": getattr( + _metadata_raw, "user_api_key_alias", None + ), + "user_api_key_team_id": getattr( + _metadata_raw, "user_api_key_team_id", None + ), + "user_api_key_team_alias": getattr( + _metadata_raw, "user_api_key_team_alias", None + ), + "user_api_key_hash": getattr(_metadata_raw, "user_api_key_hash", None), + "requester_ip_address": getattr( + _metadata_raw, "requester_ip_address", None + ), + "user_agent": getattr(_metadata_raw, "user_agent", None), + } + _litellm_params_metadata = _litellm_params.get("metadata", {}) or {} + + # Extract user_api_key_auth if present (proxy injects this, skipped in merge) + user_api_key_auth = _litellm_params_metadata.get("user_api_key_auth") + + def _get_api_key_alias() -> Optional[str]: + val = _metadata.get("user_api_key_alias") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_alias") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "key_alias", None) + return None + + def _get_team_id() -> Optional[str]: + val = _metadata.get("user_api_key_team_id") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_team_id") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "team_id", None) + return None + + def _get_team_alias() -> Optional[str]: + val = _metadata.get("user_api_key_team_alias") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_team_alias") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "team_alias", None) + return None + + def _get_hashed_api_key() -> Optional[str]: + val = _metadata.get("user_api_key_hash") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_hash") + if val is not None: + return val + if user_api_key_auth is not None: + return getattr(user_api_key_auth, "api_key", None) or getattr( + user_api_key_auth, "api_key_hash", None + ) + return None + + return { + "api_key_alias": _get_api_key_alias(), + "team": _get_team_id(), + "team_alias": _get_team_alias(), + "hashed_api_key": _get_hashed_api_key(), + "client_ip": _metadata.get("requester_ip_address") + or _litellm_params_metadata.get("requester_ip_address"), + "user_agent": _metadata.get("user_agent") + or _litellm_params_metadata.get("user_agent"), + } + def set_llm_deployment_failure_metrics(self, request_kwargs: dict): """ Sets Failure metrics when an LLM API call fails @@ -1707,6 +1809,21 @@ class PrometheusLogger(CustomLogger): model_id = standard_logging_payload.get("model_id", None) exception = request_kwargs.get("exception", None) + # Fallback: model_id from litellm_metadata.model_info + if model_id is None: + _model_info = ( + (_litellm_params.get("litellm_metadata") or {}).get("model_info") + or (_litellm_params.get("metadata") or {}).get("model_info") + or {} + ) + model_id = _model_info.get("id") + + # Fallback: model_group from litellm_metadata + if model_group is None: + model_group = (_litellm_params.get("litellm_metadata") or {}).get( + "model_group" + ) or (_litellm_params.get("metadata") or {}).get("model_group") + llm_provider = _litellm_params.get("custom_llm_provider", None) if self._should_skip_metrics_for_invalid_key( @@ -1714,9 +1831,37 @@ class PrometheusLogger(CustomLogger): standard_logging_payload=standard_logging_payload, ): return - hashed_api_key = standard_logging_payload.get("metadata", {}).get( + + # Extract context labels from all available sources (fix for None labels) + fallback_values = self._extract_deployment_failure_label_values( + request_kwargs + ) + _metadata = standard_logging_payload.get("metadata", {}) or {} + hashed_api_key = fallback_values.get("hashed_api_key") or _metadata.get( "user_api_key_hash" ) + api_key_alias = fallback_values.get("api_key_alias") or _metadata.get( + "user_api_key_alias" + ) + team = fallback_values.get("team") or _metadata.get("user_api_key_team_id") + team_alias = fallback_values.get("team_alias") or _metadata.get( + "user_api_key_team_alias" + ) + client_ip = fallback_values.get("client_ip") or _metadata.get( + "requester_ip_address" + ) + user_agent = fallback_values.get("user_agent") or _metadata.get( + "user_agent" + ) + + # exception_status: prefer status_code, fallback to exception class for known types + exception_status = None + if exception is not None: + exception_status = str(getattr(exception, "status_code", None)) + if exception_status == "None" or not exception_status: + code = getattr(exception, "code", None) + if code is not None: + exception_status = str(code) # Create enum_values for the label factory (always create for use in different metrics) enum_values = UserAPIKeyLabelValues( @@ -1724,26 +1869,18 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, api_provider=llm_provider, - exception_status=( - str(getattr(exception, "status_code", None)) if exception else None - ), + exception_status=exception_status, exception_class=( self._get_exception_class_name(exception) if exception else None ), - requested_model=model_group, + requested_model=model_group or litellm_model_name, hashed_api_key=hashed_api_key, - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], + api_key_alias=api_key_alias, + team=team, + team_alias=team_alias, tags=standard_logging_payload.get("request_tags", []), - client_ip=standard_logging_payload["metadata"].get( - "requester_ip_address" - ), - user_agent=standard_logging_payload["metadata"].get("user_agent"), + client_ip=client_ip, + user_agent=user_agent, ) """ 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 0a57d046c7..c39454728a 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 @@ -1,4 +1,3 @@ -import io import os import sys @@ -10,13 +9,10 @@ from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, call, patch import pytest -from prometheus_client import REGISTRY, CollectorRegistry +from prometheus_client import REGISTRY import litellm -from litellm import completion from litellm._logging import verbose_logger -from litellm._uuid import uuid -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -37,7 +33,6 @@ from litellm.proxy._types import UserAPIKeyAuth verbose_logger.setLevel(logging.DEBUG) litellm.set_verbose = True -import time @pytest.fixture @@ -293,7 +288,6 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): ) as mock_get_team, patch( "litellm.proxy.auth.auth_checks.get_key_object" ) as mock_get_key: - mock_get_team.return_value = MagicMock(budget_reset_at=future_reset_time_team) mock_get_key.return_value = MagicMock(budget_reset_at=future_reset_time_key) @@ -648,25 +642,16 @@ async def test_async_log_failure_event(prometheus_logger): ) # litellm_llm_api_failed_requests_metric incremented - """ - Expected metrics - end_user_id, - user_api_key, - user_api_key_alias, - model, - user_api_team, - user_api_team_alias, - user_id, - """ + # Labels: end_user, api_key_hash, api_key_alias, model, team, team_alias, user, model_id prometheus_logger.litellm_llm_api_failed_requests_metric.labels.assert_called_once_with( - None, + None, # end_user_id "test_hash", "test_alias", "gpt-3.5-turbo", "test_team", "test_team_alias", "test_user", - "model-123", + "model-123", # model_id from standard_logging_payload ) prometheus_logger.litellm_llm_api_failed_requests_metric.labels().inc.assert_called_once() @@ -678,38 +663,54 @@ async def test_async_log_failure_event(prometheus_logger): api_provider="openai", ) - # deployment failure responses incremented - prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once_with( - litellm_model_name="gpt-3.5-turbo", - model_id="model-123", - api_base="https://api.openai.com", - api_provider="openai", - exception_status="None", - exception_class="Exception", - requested_model="openai-gpt", # passed in standard logging payload - hashed_api_key="test_hash", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - client_ip="127.0.0.1", # from standard logging payload - user_agent=None, + # deployment failure responses incremented - verify key labels are populated + prometheus_logger.litellm_deployment_failure_responses.labels.assert_called_once() + actual_failure_labels = ( + prometheus_logger.litellm_deployment_failure_responses.labels.call_args.kwargs ) + expected_failure_labels = { + "litellm_model_name": "gpt-3.5-turbo", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "api_provider": "openai", + "exception_class": "Exception", + "requested_model": "openai-gpt", + "hashed_api_key": "test_hash", + "api_key_alias": "test_alias", + "team": "test_team", + "team_alias": "test_team_alias", + } + for key, expected_val in expected_failure_labels.items(): + assert key in actual_failure_labels, f"Missing label {key}" + assert ( + actual_failure_labels[key] == expected_val + ), f"Label {key}: expected {expected_val!r}, got {actual_failure_labels[key]!r}" + assert actual_failure_labels.get("exception_status") in ("None", None) + assert actual_failure_labels.get("client_ip") == "127.0.0.1" prometheus_logger.litellm_deployment_failure_responses.labels().inc.assert_called_once() - # deployment total requests incremented - prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once_with( - litellm_model_name="gpt-3.5-turbo", - model_id="model-123", - api_base="https://api.openai.com", - api_provider="openai", - requested_model="openai-gpt", # passed in standard logging payload - hashed_api_key="test_hash", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - client_ip="127.0.0.1", # from standard logging payload - user_agent=None, + # deployment total requests incremented - verify key labels are populated + prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once() + actual_total_labels = ( + prometheus_logger.litellm_deployment_total_requests.labels.call_args.kwargs ) + expected_total_labels = { + "litellm_model_name": "gpt-3.5-turbo", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "api_provider": "openai", + "requested_model": "openai-gpt", + "hashed_api_key": "test_hash", + "api_key_alias": "test_alias", + "team": "test_team", + "team_alias": "test_team_alias", + } + for key, expected_val in expected_total_labels.items(): + assert key in actual_total_labels, f"Missing label {key}" + assert ( + actual_total_labels[key] == expected_val + ), f"Label {key}: expected {expected_val!r}, got {actual_total_labels[key]!r}" + assert actual_total_labels.get("client_ip") == "127.0.0.1" prometheus_logger.litellm_deployment_total_requests.labels().inc.assert_called_once() @@ -1095,7 +1096,7 @@ def test_increment_deployment_cooled_down(prometheus_logger): import inspect method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down) - expected_label_count = len([p for p in method_sig.parameters.keys() if p != 'self']) + expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"]) mock_chain = MagicMock() @@ -1103,11 +1104,15 @@ def test_increment_deployment_cooled_down(prometheus_logger): """Validate label count matches metric definition""" total = len(label_values) + len(label_kwargs) if total != expected_label_count: - raise ValueError(f"Incorrect label count: expected {expected_label_count}, got {total}") + raise ValueError( + f"Incorrect label count: expected {expected_label_count}, got {total}" + ) return mock_chain prometheus_logger.litellm_deployment_cooled_down = MagicMock() - prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(side_effect=validating_labels) + prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock( + side_effect=validating_labels + ) prometheus_logger.increment_deployment_cooled_down( litellm_model_name="gpt-3.5-turbo", @@ -1179,8 +1184,12 @@ def test_get_custom_labels_from_top_level_metadata(monkeypatch): metadata = { "requester_ip_address": "10.48.203.20", # Top-level field "user_api_key_alias": "TestAlias", # Top-level field - "requester_metadata": {"nested_field": "nested_value"}, # Nested dict (excluded) - "user_api_key_auth_metadata": {"another_nested": "value"}, # Nested dict (excluded) + "requester_metadata": { + "nested_field": "nested_value" + }, # Nested dict (excluded) + "user_api_key_auth_metadata": { + "another_nested": "value" + }, # Nested dict (excluded) } result = get_custom_labels_from_metadata(metadata) assert result == { @@ -1217,7 +1226,9 @@ def test_get_custom_labels_from_top_level_and_nested_metadata(monkeypatch): } -async def test_async_log_success_event_with_top_level_metadata(prometheus_logger, monkeypatch): +async def test_async_log_success_event_with_top_level_metadata( + prometheus_logger, monkeypatch +): """ Test that async_log_success_event correctly extracts custom labels from top-level metadata fields like requester_ip_address, not just from nested dictionaries. @@ -1231,7 +1242,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger standard_logging_object = create_standard_logging_payload() standard_logging_object["metadata"]["requester_ip_address"] = "10.48.203.20" standard_logging_object["metadata"]["requester_metadata"] = {} # Empty nested dict - standard_logging_object["metadata"]["user_api_key_auth_metadata"] = {} # Empty nested dict + standard_logging_object["metadata"][ + "user_api_key_auth_metadata" + ] = {} # Empty nested dict kwargs = { "model": "gpt-3.5-turbo", @@ -1273,7 +1286,9 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger prometheus_logger.litellm_remaining_user_budget_metric = create_mock_metric() prometheus_logger.litellm_user_max_budget_metric = create_mock_metric() prometheus_logger.litellm_user_budget_remaining_hours_metric = create_mock_metric() - prometheus_logger.litellm_remaining_api_key_requests_for_model = create_mock_metric() + prometheus_logger.litellm_remaining_api_key_requests_for_model = ( + create_mock_metric() + ) prometheus_logger.litellm_remaining_api_key_tokens_for_model = create_mock_metric() prometheus_logger.litellm_llm_api_time_to_first_token_metric = create_mock_metric() prometheus_logger.litellm_llm_api_latency_metric = create_mock_metric() @@ -1302,7 +1317,7 @@ async def test_async_log_success_event_with_top_level_metadata(prometheus_logger # This confirms that the custom label extraction logic ran without errors assert prometheus_logger.litellm_requests_metric.labels.called assert prometheus_logger.litellm_spend_metric.labels.called - + # Verify that the labels() method was called with some arguments (either positional or keyword) # This ensures the custom label extraction happened and didn't cause a "Incorrect label names" error call_args = prometheus_logger.litellm_requests_metric.labels.call_args @@ -1494,7 +1509,6 @@ async def test_initialize_remaining_budget_metrics(prometheus_logger): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.team_endpoints.get_paginated_teams" ) as mock_get_teams: - # Create mock team data with proper datetime objects for budget_reset_at future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now mock_teams = [ @@ -1592,21 +1606,22 @@ async def test_initialize_remaining_budget_metrics_exception_handling( ) as mock_get_teams, patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: - # Make get_paginated_teams raise an exception mock_get_teams.side_effect = Exception("Database error") mock_list_keys.side_effect = Exception("Key listing error") - + # Mock prisma_client structure to raise an exception for user budget metrics # The code accesses prisma_client.db.litellm_usertable.find_many and count mock_usertable = MagicMock() - mock_usertable.find_many = MagicMock(side_effect=Exception("User database error")) + mock_usertable.find_many = MagicMock( + side_effect=Exception("User database error") + ) mock_usertable.count = MagicMock(side_effect=Exception("User count error")) - + # Mock litellm_teamtable to raise an exception for team count metrics mock_teamtable = MagicMock() mock_teamtable.count = MagicMock(side_effect=Exception("Team count error")) - + mock_db = MagicMock() mock_db.litellm_usertable = mock_usertable mock_db.litellm_teamtable = mock_teamtable @@ -1661,7 +1676,6 @@ async def test_initialize_api_key_budget_metrics(prometheus_logger): with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: - # Create mock key data with proper datetime objects for budget_reset_at future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now key1 = UserAPIKeyAuth( @@ -1916,7 +1930,6 @@ def test_prometheus_label_factory_with_custom_tags(monkeypatch): Test that prometheus_label_factory correctly handles custom tags """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -1954,7 +1967,6 @@ def test_prometheus_label_factory_with_no_custom_tags(monkeypatch): Test that prometheus_label_factory works when no custom tags are configured """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -2179,9 +2191,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): All three metrics should be properly incremented when making a successful completion request. """ - from prometheus_client import CollectorRegistry, Counter - import litellm from litellm.types.integrations.prometheus import PrometheusMetricsConfig # Clear registry before test