scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(prometheus): configure deployment caller identity

Browse source
This commit is contained in:
Mark Philipp 2026-08-25 09:39:30 -07:00
parent bb27bfd9a7
commit 8d0c23ea64
5 changed files with 475 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/__init__.py
View file

@ -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

15
litellm/integrations/prometheus.py
View file

@ -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", []),

7
litellm/proxy/proxy_server.py
View file

@ -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"

58
litellm/types/integrations/prometheus.py
View file

@ -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

392
tests/test_litellm/integrations/test_prometheus_caller_identity.py Normal file
View file

@ -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"