From 8d0c23ea6491b37b7d02edcb34e8a722aece4156 Mon Sep 17 00:00:00 2001 From: Mark Philipp Date: Tue, 25 Aug 2026 09:39:30 -0700 Subject: [PATCH] feat(prometheus): configure deployment caller identity --- litellm/__init__.py | 5 + litellm/integrations/prometheus.py | 15 + litellm/proxy/proxy_server.py | 7 + litellm/types/integrations/prometheus.py | 58 ++- .../test_prometheus_caller_identity.py | 392 ++++++++++++++++++ 5 files changed, 475 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_caller_identity.py diff --git a/litellm/__init__.py b/litellm/__init__.py index e95b553c5d4..3e77f621567 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -463,6 +463,11 @@ prometheus_metrics_config: Optional[List] = None prometheus_exclude_metrics: Optional[List[str]] = None prometheus_exclude_labels: Optional[List[str]] = None prometheus_emit_stream_label: bool = False +prometheus_deployment_and_latency_caller_identity: Literal[ + "api_key_alias", + "user_email", + "both", +] = "api_key_alias" # Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on # `litellm_proxy_failed_requests_metric`. Off by default to preserve the # pre-unification label set so existing dashboards / recording rules keyed on diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f9195db1d67..2b0a0658650 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2462,6 +2462,7 @@ class PrometheusLogger(CustomLogger): else: _metadata = { "user_api_key_alias": getattr(_metadata_raw, "user_api_key_alias", None), + "user_api_key_user_email": getattr(_metadata_raw, "user_api_key_user_email", 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), @@ -2484,6 +2485,17 @@ class PrometheusLogger(CustomLogger): return getattr(user_api_key_auth, "key_alias", None) return None + def _get_user_email() -> str | None: + val = _metadata.get("user_api_key_user_email") + if val is not None: + return val + val = _litellm_params_metadata.get("user_api_key_user_email") + if val is not None: + return val + if user_api_key_auth is not None: + return self._safe_get(user_api_key_auth, "user_email") + return None + def _get_team_id() -> str | None: val = _metadata.get("user_api_key_team_id") if val is not None: @@ -2519,6 +2531,7 @@ class PrometheusLogger(CustomLogger): return { "api_key_alias": _get_api_key_alias(), + "user_email": _get_user_email(), "team": _get_team_id(), "team_alias": _get_team_alias(), "hashed_api_key": _get_hashed_api_key(), @@ -2576,6 +2589,7 @@ class PrometheusLogger(CustomLogger): _metadata: Final = standard_logging_payload.get("metadata", {}) or {} hashed_api_key: Final = fallback_values.get("hashed_api_key") or _metadata.get("user_api_key_hash") api_key_alias: Final = fallback_values.get("api_key_alias") or _metadata.get("user_api_key_alias") + user_email: Final = fallback_values.get("user_email") team: Final = fallback_values.get("team") or _metadata.get("user_api_key_team_id") team_alias: Final = fallback_values.get("team_alias") or _metadata.get("user_api_key_team_alias") client_ip: Final = fallback_values.get("client_ip") or _metadata.get("requester_ip_address") @@ -2616,6 +2630,7 @@ class PrometheusLogger(CustomLogger): requested_model=label_requested_model, hashed_api_key=hashed_api_key, api_key_alias=api_key_alias, + user_email=user_email, team=team, team_alias=team_alias, tags=standard_logging_payload.get("request_tags", []), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dced4e26b6..8925e701b97 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4852,6 +4852,13 @@ class ProxyConfig: if litellm_settings is None: litellm_settings = {} if litellm_settings: + # Prometheus collectors have fixed label schemas. Load this setting + # before processing callbacks so YAML key order cannot construct the + # collectors with the default caller-identity mode. + caller_identity_mode: Final = litellm_settings.get("prometheus_deployment_and_latency_caller_identity") + if caller_identity_mode is not None: + litellm.prometheus_deployment_and_latency_caller_identity = caller_identity_mode + # ANSI escape code for blue text blue_color_code: Final = "\033[94m" reset_color_code: Final = "\033[0m" diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index ebec5df55fa..dbd66ffaf5e 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,5 +1,5 @@ import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import MISSING, dataclass, field, fields from enum import Enum from types import MappingProxyType @@ -276,6 +276,57 @@ DEFINED_PROMETHEUS_METRICS = Literal[ ] +PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_METRICS: Final[frozenset[str]] = frozenset( + { + "litellm_deployment_total_requests", + "litellm_deployment_success_responses", + "litellm_deployment_failure_responses", + "litellm_request_total_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_overhead_latency_metric", + "litellm_deployment_latency_per_output_token", + } +) + +PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_VALUES: Final[tuple[str, ...]] = ( + "api_key_alias", + "user_email", + "both", +) + + +def _resolve_deployment_and_latency_caller_identity_labels( + metric_name: str, + labels: Sequence[object], +) -> list[str]: # mutable-ok: every caller must receive an independently mutable label list + """Return a fresh label list with the configured caller identity schema.""" + if not all(isinstance(label, str) for label in labels): + raise TypeError(f"Prometheus labels for {metric_name} must be strings") + resolved_labels: Final = [label for label in labels if isinstance(label, str)] + if metric_name not in PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_METRICS: + return resolved_labels + + caller_identity: Final[object] = getattr( + litellm, + "prometheus_deployment_and_latency_caller_identity", + "api_key_alias", + ) + if caller_identity not in PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_VALUES: + accepted_values: Final = ", ".join(PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_VALUES) + raise ValueError( + "Invalid prometheus_deployment_and_latency_caller_identity=" + f"{caller_identity!r}. Accepted values: {accepted_values}." + ) + + alias_index: Final = resolved_labels.index(UserAPIKeyLabelNames.API_KEY_ALIAS.value) + if caller_identity == "user_email": + resolved_labels[alias_index] = UserAPIKeyLabelNames.USER_EMAIL.value + elif caller_identity == "both": + resolved_labels.insert(alias_index + 1, UserAPIKeyLabelNames.USER_EMAIL.value) + + return resolved_labels + + class PrometheusMetricLabels: litellm_llm_api_latency_metric = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -781,7 +832,10 @@ class PrometheusMetricLabels: @staticmethod def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> list[str]: - default_labels: Final = getattr(PrometheusMetricLabels, label_name) + default_labels: Final = _resolve_deployment_and_latency_caller_identity_labels( + metric_name=label_name, + labels=getattr(PrometheusMetricLabels, label_name), + ) custom_labels: Final = [] # Add custom metadata labels diff --git a/tests/test_litellm/integrations/test_prometheus_caller_identity.py b/tests/test_litellm/integrations/test_prometheus_caller_identity.py new file mode 100644 index 00000000000..b5b65049ff8 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_caller_identity.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +from datetime import datetime, timedelta +from pathlib import Path +from typing import Final, cast +from unittest.mock import patch + +import pytest +import yaml +from prometheus_client import REGISTRY, generate_latest +from prometheus_client.parser import text_string_to_metric_families + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.integrations.prometheus import ( + DEFINED_PROMETHEUS_METRICS, + PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_METRICS, + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, +) +from litellm.types.utils import StandardLoggingPayload + +TARGET_METRICS: Final[tuple[DEFINED_PROMETHEUS_METRICS, ...]] = cast( + tuple[DEFINED_PROMETHEUS_METRICS, ...], + tuple(sorted(PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_METRICS)), +) +IDENTITY_MODES: Final = ("api_key_alias", "user_email", "both") + + +def _clear_prometheus_registry() -> None: + for collector in list(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] + REGISTRY.unregister(collector) + + +@pytest.fixture(autouse=True) +def reset_prometheus_settings(monkeypatch: pytest.MonkeyPatch): + _clear_prometheus_registry() + monkeypatch.setattr(litellm, "prometheus_deployment_and_latency_caller_identity", "api_key_alias") + monkeypatch.setattr(litellm, "prometheus_metrics_config", None) + monkeypatch.setattr(litellm, "prometheus_exclude_metrics", None) + monkeypatch.setattr(litellm, "prometheus_exclude_labels", None) + monkeypatch.setattr(litellm, "custom_prometheus_metadata_labels", []) + monkeypatch.setattr(litellm, "custom_prometheus_tags", []) + yield + _clear_prometheus_registry() + + +def _expected_identity_labels(baseline: list[str], mode: str) -> list[str]: + expected = list(baseline) + alias_index = expected.index(UserAPIKeyLabelNames.API_KEY_ALIAS.value) + if mode == "user_email": + expected[alias_index] = UserAPIKeyLabelNames.USER_EMAIL.value + elif mode == "both": + expected.insert(alias_index + 1, UserAPIKeyLabelNames.USER_EMAIL.value) + return expected + + +def _set_caller_identity(monkeypatch: pytest.MonkeyPatch, mode: str) -> None: + monkeypatch.setattr(litellm, "prometheus_deployment_and_latency_caller_identity", mode) + + +@pytest.mark.parametrize("metric_name", TARGET_METRICS) +@pytest.mark.parametrize("mode", IDENTITY_MODES) +def test_target_metric_label_schema_for_each_caller_identity_mode( + monkeypatch: pytest.MonkeyPatch, + metric_name: DEFINED_PROMETHEUS_METRICS, + mode: str, +): + _set_caller_identity(monkeypatch, "api_key_alias") + baseline = PrometheusMetricLabels.get_labels(metric_name) + + _set_caller_identity(monkeypatch, mode) + actual = PrometheusMetricLabels.get_labels(metric_name) + + assert actual == _expected_identity_labels(baseline, mode) + + +def test_repeated_label_resolution_does_not_mutate_class_level_or_shared_lists( + monkeypatch: pytest.MonkeyPatch, +): + total_request_labels = PrometheusMetricLabels.litellm_deployment_total_requests + success_labels = PrometheusMetricLabels.litellm_deployment_success_responses + original = tuple(total_request_labels) + + assert success_labels is total_request_labels + for mode in (*IDENTITY_MODES, *reversed(IDENTITY_MODES)): + _set_caller_identity(monkeypatch, mode) + for metric_name in TARGET_METRICS: + resolved = PrometheusMetricLabels.get_labels(metric_name) + assert resolved is not getattr(PrometheusMetricLabels, metric_name) + + assert PrometheusMetricLabels.litellm_deployment_total_requests is total_request_labels + assert PrometheusMetricLabels.litellm_deployment_success_responses is success_labels + assert success_labels is total_request_labels + assert tuple(total_request_labels) == original + + +def test_invalid_caller_identity_mode_fails_during_prometheus_initialization( + monkeypatch: pytest.MonkeyPatch, +): + _set_caller_identity(monkeypatch, "invalid") + + with pytest.raises(ValueError) as exc_info: + PrometheusLogger() + + message = str(exc_info.value) + assert "prometheus_deployment_and_latency_caller_identity" in message + for accepted_value in IDENTITY_MODES: + assert accepted_value in message + + +@pytest.mark.parametrize( + ("mode", "include_labels", "is_valid"), + ( + ("api_key_alias", ["api_key_alias"], True), + ("api_key_alias", ["user_email"], False), + ("user_email", ["user_email"], True), + ("user_email", ["api_key_alias"], False), + ("both", ["api_key_alias"], True), + ("both", ["user_email"], True), + ("both", ["api_key_alias", "user_email"], True), + ), +) +def test_include_labels_validation_matches_caller_identity_mode( + monkeypatch: pytest.MonkeyPatch, + mode: str, + include_labels: list[str], + is_valid: bool, +): + _set_caller_identity(monkeypatch, mode) + monkeypatch.setattr( + litellm, + "prometheus_metrics_config", + [ + { + "group": "caller_identity", + "metrics": ["litellm_deployment_total_requests"], + "include_labels": include_labels, + } + ], + ) + + if not is_valid: + with pytest.raises(ValueError, match="Configuration validation failed"): + PrometheusLogger() + return + + logger = PrometheusLogger() + assert logger.get_labels_for_metric("litellm_deployment_total_requests") == include_labels + + +@pytest.mark.parametrize( + ("mode", "exclude_labels", "remaining_identity_labels"), + ( + ("api_key_alias", ["api_key_alias"], set[str]()), + ("api_key_alias", ["user_email"], {"api_key_alias"}), + ("user_email", ["user_email"], set[str]()), + ("user_email", ["api_key_alias"], {"user_email"}), + ("both", ["api_key_alias"], {"user_email"}), + ("both", ["user_email"], {"api_key_alias"}), + ("both", ["api_key_alias", "user_email"], set[str]()), + ), +) +def test_exclude_labels_can_remove_supported_identity_labels( + monkeypatch: pytest.MonkeyPatch, + mode: str, + exclude_labels: list[str], + remaining_identity_labels: set[str], +): + _set_caller_identity(monkeypatch, mode) + monkeypatch.setattr(litellm, "prometheus_exclude_labels", exclude_labels) + + logger = PrometheusLogger() + labels = logger.get_labels_for_metric("litellm_deployment_total_requests") + + assert set(labels) & {"api_key_alias", "user_email"} == remaining_identity_labels + + +@pytest.mark.parametrize("mode", IDENTITY_MODES) +def test_non_target_metric_label_schema_is_unchanged(monkeypatch: pytest.MonkeyPatch, mode: str): + baseline = list(PrometheusMetricLabels.litellm_overhead_with_guardrails_latency_metric) + _set_caller_identity(monkeypatch, mode) + + actual = PrometheusMetricLabels.get_labels("litellm_overhead_with_guardrails_latency_metric") + + assert actual == baseline + assert "api_key_alias" in actual + assert "user_email" not in actual + + +def _standard_logging_payload(user_email: str | None = "alice@example.com") -> StandardLoggingPayload: + return cast( + StandardLoggingPayload, + { + "api_base": "https://api.example.com", + "model_group": "requested-model", + "model_id": "deployment-id", + "request_tags": [], + "metadata": { + "user_api_key_hash": "hashed-key", + "user_api_key_alias": "alias-a", + "user_api_key_user_email": user_email, + "user_api_key_team_id": "team-id", + "user_api_key_team_alias": "team-alias", + "requester_ip_address": "192.0.2.10", + "user_agent": "caller-identity-test", + }, + "hidden_params": { + "additional_headers": None, + "litellm_overhead_time_ms": 125, + }, + }, + ) + + +def _sample_labels(scrape: str, sample_name: str) -> list[dict[str, str]]: + return [ + sample.labels + for family in text_string_to_metric_families(scrape) + for sample in family.samples + if sample.name == sample_name + ] + + +@pytest.mark.parametrize("mode", IDENTITY_MODES) +def test_successful_request_emits_configured_identity_on_real_counter_and_histogram_samples( + monkeypatch: pytest.MonkeyPatch, + mode: str, +): + _set_caller_identity(monkeypatch, mode) + logger = PrometheusLogger() + payload = _standard_logging_payload() + enum_values = UserAPIKeyLabelValues( + end_user="end-user", + user="user-id", + user_email="alice@example.com", + hashed_api_key="hashed-key", + api_key_alias="alias-a", + requested_model="requested-model", + model_group="requested-model", + team="team-id", + team_alias="team-alias", + model="provider-model", + litellm_model_name="deployment-model", + model_id="deployment-id", + api_base="https://api.example.com", + api_provider="openai", + client_ip="192.0.2.10", + user_agent="caller-identity-test", + ) + start_time = datetime.now() + api_call_start_time = start_time + timedelta(milliseconds=100) + completion_start_time = api_call_start_time + timedelta(milliseconds=200) + end_time = start_time + timedelta(seconds=1) + request_kwargs = { + "model": "deployment-model", + "stream": True, + "start_time": start_time, + "api_call_start_time": api_call_start_time, + "completion_start_time": completion_start_time, + "end_time": end_time, + "litellm_params": { + "custom_llm_provider": "openai", + "metadata": {"model_info": {"id": "deployment-id"}}, + }, + "standard_logging_object": payload, + } + + logger._set_latency_metrics( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] + kwargs=request_kwargs, + model="deployment-model", + user_api_key="hashed-key", + user_api_key_alias="alias-a", + user_api_team="team-id", + user_api_team_alias="team-alias", + enum_values=enum_values, + ) + logger.set_llm_deployment_success_metrics( # pyright: ignore[reportUnknownMemberType] + request_kwargs=request_kwargs, + start_time=start_time, + end_time=end_time, + enum_values=enum_values, + output_tokens=10, + ) + + scrape = generate_latest(REGISTRY).decode() + sample_names = ( + "litellm_deployment_total_requests_total", + "litellm_deployment_success_responses_total", + "litellm_request_total_latency_metric_count", + "litellm_llm_api_time_to_first_token_metric_count", + "litellm_overhead_latency_metric_count", + "litellm_deployment_latency_per_output_token_count", + ) + for sample_name in sample_names: + samples = _sample_labels(scrape, sample_name) + assert len(samples) == 1, sample_name + labels = samples[0] + if mode == "api_key_alias": + assert labels["api_key_alias"] == "alias-a" + assert "user_email" not in labels + elif mode == "user_email": + assert labels["user_email"] == "alice@example.com" + assert "api_key_alias" not in labels + else: + assert labels["api_key_alias"] == "alias-a" + assert labels["user_email"] == "alice@example.com" + + +@pytest.mark.parametrize( + ("standard_email", "metadata_email", "auth_email", "expected_email"), + ( + ("standard@example.com", "metadata@example.com", "auth@example.com", "standard@example.com"), + (None, "metadata@example.com", "auth@example.com", "metadata@example.com"), + (None, None, "auth@example.com", "auth@example.com"), + (None, None, None, "None"), + ), +) +def test_deployment_failure_email_fallbacks_reach_both_real_counters( + monkeypatch: pytest.MonkeyPatch, + standard_email: str | None, + metadata_email: str | None, + auth_email: str | None, + expected_email: str, +): + _set_caller_identity(monkeypatch, "both") + logger = PrometheusLogger() + payload = _standard_logging_payload(user_email=standard_email) + metadata = { + "model_info": {"id": "deployment-id"}, + "user_api_key_user_email": metadata_email, + "user_api_key_auth": UserAPIKeyAuth(user_email=auth_email), + } + request_kwargs = { + "model": "deployment-model", + "litellm_params": { + "custom_llm_provider": "openai", + "metadata": metadata, + }, + "standard_logging_object": payload, + "exception": RuntimeError("provider failed"), + } + + logger.set_llm_deployment_failure_metrics(request_kwargs) # pyright: ignore[reportUnknownMemberType] + + scrape = generate_latest(REGISTRY).decode() + for sample_name in ( + "litellm_deployment_failure_responses_total", + "litellm_deployment_total_requests_total", + ): + samples = _sample_labels(scrape, sample_name) + assert len(samples) == 1, sample_name + assert samples[0]["api_key_alias"] == "alias-a" + assert samples[0]["user_email"] == expected_email + + +@pytest.mark.asyncio +async def test_proxy_config_loads_caller_identity_before_initializing_callbacks(tmp_path: Path): + from litellm.proxy.proxy_server import ProxyConfig + + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-4", "api_key": "test-key"}, + } + ], + "litellm_settings": { + "callbacks": ["prometheus"], + "prometheus_deployment_and_latency_caller_identity": "both", + }, + }, + sort_keys=False, + ) + ) + observed_modes: list[str] = [] + + def capture_mode(*args: object, **kwargs: object) -> None: + observed_modes.append(litellm.prometheus_deployment_and_latency_caller_identity) + + with patch( # test-quality-ok: callback interception verifies schema selection before construction + "litellm.proxy.proxy_server.initialize_callbacks_on_proxy", side_effect=capture_mode + ): + await ProxyConfig().load_config(router=None, config_file_path=str(config_path)) + + assert observed_modes == ["both"] + assert litellm.prometheus_deployment_and_latency_caller_identity == "both"