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

fixes failure metrics labels (#20152)

Browse Source 
Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
This commit is contained in:
Harshit Jain
2026-02-02 19:20:14 -08:00
committed by GitHub
co-authored by Krish Dholakia
parent 76a399ba69
commit 04c348e7bb
2 changed files with 230 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/prometheus.py
+153 -16
View File
@@ -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,
)
"""
tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
+77 -67
View File
@@ -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