scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83
litellm/tests/test_litellm/integrations/test_prometheus_caller_identity.py
mphilippnv e52f05566d
feat(prometheus): configure deployment caller identity (#38221) 
* feat(prometheus): configure deployment caller identity

* test(prometheus): satisfy strict caller identity lint

* fix(prometheus): align caller identity on latency metrics

* fix(prometheus): validate caller identity mode before collectors register

Fail config load on an invalid prometheus_deployment_and_latency_caller_identity
value (including null) and on include_labels entries the selected mode removes
from a target metric, instead of booting green with an empty /metrics.
Validate the mode at the top of PrometheusLogger.__init__ so an invalid value
raises before any collector lands in the process-global registry, keeping
retries free of duplicated-timeseries errors. Label-validation errors now name
the mode setting alongside the rejected label.

---------

Co-authored-by: Mark Philipp <mphilipp622@gmail.com>
Co-authored-by: Yucheng Zhu <yucheng@berri.ai>
2026-08-25 23:06:02 -07:00

411 lines
15 KiB
Python
Raw Permalink Blame History

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,
match="prometheus_deployment_and_latency_caller_identity",
) 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
def test_label_resolution_rejects_non_string_class_labels(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
PrometheusMetricLabels,
"litellm_deployment_total_requests",
["api_key_alias", 1],
)
with pytest.raises(TypeError, match=r"Prometheus labels .* must be strings"):
PrometheusMetricLabels.get_labels("litellm_deployment_total_requests")
@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"},
"queue_time_seconds": 0.05,
},
},
"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_latency_metric_count",
"litellm_llm_api_time_to_first_token_metric_count",
"litellm_request_queue_time_seconds_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"
Reference in a new issue View git blame Copy permalink