litellm/tests/unit/integrations/test_prometheus_caller_identity.py

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,
    LabelValidationError,
    PrometheusMetricLabels,
    UserAPIKeyLabelNames,
    UserAPIKeyLabelValues,
    validate_caller_identity_settings,
    validate_prometheus_deployment_and_latency_caller_identity,
)
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 = _write_proxy_config(
        tmp_path,
        {
            "callbacks": ["prometheus"],
            "prometheus_deployment_and_latency_caller_identity": "both",
        },
    )
    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"


def _identity_settings(mode: object, metrics_config: object = None) -> dict[str, object]:
    settings: dict[str, object] = {"prometheus_deployment_and_latency_caller_identity": mode}
    if metrics_config is not None:
        settings["prometheus_metrics_config"] = metrics_config
    return settings


def test_validate_mode_returns_each_accepted_value_and_defaults_to_api_key_alias(
    monkeypatch: pytest.MonkeyPatch,
):
    for mode in IDENTITY_MODES:
        _set_caller_identity(monkeypatch, mode)
        assert validate_prometheus_deployment_and_latency_caller_identity() == mode

    monkeypatch.delattr(litellm, "prometheus_deployment_and_latency_caller_identity")
    assert validate_prometheus_deployment_and_latency_caller_identity() == "api_key_alias"


def test_accepted_values_constant_matches_parametrized_modes():
    from litellm.types.integrations.prometheus import (
        PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_VALUES,
    )

    assert PROMETHEUS_DEPLOYMENT_AND_LATENCY_CALLER_IDENTITY_VALUES == IDENTITY_MODES
    assert len(TARGET_METRICS) == 9


@pytest.mark.parametrize(
    "invalid_mode",
    ("user-email", "USER_EMAIL", "", None, True, 1, ["user_email"], {"mode": "user_email"}),
)
def test_validate_mode_rejects_invalid_values_and_names_accepted_ones(
    monkeypatch: pytest.MonkeyPatch,
    invalid_mode: object,
):
    monkeypatch.setattr(litellm, "prometheus_deployment_and_latency_caller_identity", invalid_mode)

    with pytest.raises(ValueError, match="prometheus_deployment_and_latency_caller_identity") as exc_info:
        validate_prometheus_deployment_and_latency_caller_identity()

    message = str(exc_info.value)
    assert repr(invalid_mode) in message
    for accepted_value in IDENTITY_MODES:
        assert accepted_value in message


def test_validate_caller_identity_settings_without_key_leaves_mode_untouched(
    monkeypatch: pytest.MonkeyPatch,
):
    _set_caller_identity(monkeypatch, "both")

    validate_caller_identity_settings({"prometheus_metrics_config": []})

    assert litellm.prometheus_deployment_and_latency_caller_identity == "both"


@pytest.mark.parametrize("mode", IDENTITY_MODES)
def test_validate_caller_identity_settings_stores_each_valid_mode(mode: str):
    validate_caller_identity_settings(_identity_settings(mode))

    assert litellm.prometheus_deployment_and_latency_caller_identity == mode


@pytest.mark.parametrize("invalid_mode", ("user-email", None))
def test_validate_caller_identity_settings_rejects_invalid_and_null_modes(invalid_mode: object):
    with pytest.raises(ValueError, match="prometheus_deployment_and_latency_caller_identity"):
        validate_caller_identity_settings(_identity_settings(invalid_mode))


def test_user_email_mode_conflict_error_names_every_conflicting_metric_and_only_those():
    metrics_config = [
        {
            "group": "non_target",
            "metrics": ["litellm_overhead_with_guardrails_latency_metric"],
            "include_labels": ["api_key_alias"],
        },
        {
            "group": "target_pair",
            "metrics": ["litellm_deployment_total_requests", "litellm_llm_api_latency_metric"],
            "include_labels": ["api_key_alias"],
        },
        {
            "group": "target_single",
            "metrics": ["litellm_request_queue_time_seconds"],
            "include_labels": ["api_key_alias"],
        },
    ]

    with pytest.raises(ValueError, match="prometheus_deployment_and_latency_caller_identity") as exc_info:
        validate_caller_identity_settings(_identity_settings("user_email", metrics_config))

    message = str(exc_info.value)
    for conflicting_metric in (
        "litellm_deployment_total_requests",
        "litellm_llm_api_latency_metric",
        "litellm_request_queue_time_seconds",
    ):
        assert conflicting_metric in message
    assert "litellm_overhead_with_guardrails_latency_metric" not in message
    assert "prometheus_deployment_and_latency_caller_identity" in message
    assert "user_email" in message


@pytest.mark.parametrize(
    ("mode", "metrics_config"),
    (
        (
            "user_email",
            [
                {
                    "group": "g",
                    "metrics": ["litellm_deployment_total_requests"],
                    "include_labels": ["user_email"],
                }
            ],
        ),
        (
            "user_email",
            [
                {
                    "group": "g",
                    "metrics": ["litellm_overhead_with_guardrails_latency_metric"],
                    "include_labels": ["api_key_alias"],
                }
            ],
        ),
        (
            "api_key_alias",
            [
                {
                    "group": "g",
                    "metrics": ["litellm_deployment_total_requests"],
                    "include_labels": ["api_key_alias"],
                }
            ],
        ),
        (
            "both",
            [
                {
                    "group": "g",
                    "metrics": ["litellm_deployment_total_requests"],
                    "include_labels": ["api_key_alias"],
                }
            ],
        ),
        ("user_email", None),
        ("user_email", ["not-a-dict"]),
        (
            "user_email",
            [{"group": "g", "metrics": ["litellm_deployment_total_requests"], "include_labels": None}],
        ),
        ("user_email", [{"group": "g", "metrics": None, "include_labels": ["api_key_alias"]}]),
    ),
)
def test_validate_caller_identity_settings_accepts_non_conflicting_configs(
    mode: str,
    metrics_config: object,
):
    settings = _identity_settings(mode)
    settings["prometheus_metrics_config"] = metrics_config

    validate_caller_identity_settings(settings)

    assert litellm.prometheus_deployment_and_latency_caller_identity == mode


def _write_proxy_config(tmp_path: Path, litellm_settings: dict[str, object]) -> Path:
    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": litellm_settings,
            },
            sort_keys=False,
        )
    )
    return config_path


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "litellm_settings",
    (
        {
            "callbacks": ["prometheus"],
            "prometheus_deployment_and_latency_caller_identity": "user-email",
        },
        {
            "callbacks": ["prometheus"],
            "prometheus_deployment_and_latency_caller_identity": None,
        },
        {
            "callbacks": ["prometheus"],
            "prometheus_deployment_and_latency_caller_identity": "user_email",
            "prometheus_metrics_config": [
                {
                    "group": "g",
                    "metrics": ["litellm_deployment_total_requests"],
                    "include_labels": ["api_key_alias"],
                }
            ],
        },
    ),
    ids=("typo-mode", "null-mode", "include-labels-conflict"),
)
async def test_proxy_config_fails_boot_before_callbacks_on_invalid_caller_identity_config(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
    litellm_settings: dict[str, object],
):
    from litellm.proxy.proxy_server import ProxyConfig

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
    monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
    config_path = _write_proxy_config(tmp_path, litellm_settings)

    with patch(  # test-quality-ok: asserts boot fails before any callback initialization
        "litellm.proxy.proxy_server.initialize_callbacks_on_proxy"
    ) as callback_init:
        with pytest.raises(ValueError, match="prometheus_deployment_and_latency_caller_identity"):
            await ProxyConfig().load_config(router=None, config_file_path=str(config_path))

    callback_init.assert_not_called()


def test_failed_init_leaves_registry_clean_so_a_corrected_retry_succeeds(
    monkeypatch: pytest.MonkeyPatch,
):
    _set_caller_identity(monkeypatch, "user-email")
    with pytest.raises(ValueError, match="prometheus_deployment_and_latency_caller_identity"):
        PrometheusLogger()

    assert list(REGISTRY._collector_to_names) == []  # pyright: ignore[reportPrivateUsage]

    _set_caller_identity(monkeypatch, "user_email")
    logger = PrometheusLogger()
    assert "user_email" in logger.get_labels_for_metric("litellm_deployment_total_requests")


@pytest.mark.parametrize("invalid_label", ("api_key_alias", "user_email"))
def test_label_validation_error_names_mode_setting_for_identity_labels_on_target_metric(
    monkeypatch: pytest.MonkeyPatch,
    invalid_label: str,
):
    _set_caller_identity(monkeypatch, "user_email")

    error = LabelValidationError(
        metric_name="litellm_deployment_total_requests",
        invalid_labels=[invalid_label],
        valid_labels=["user_email"],
    )

    assert "prometheus_deployment_and_latency_caller_identity='user_email'" in error.message
    assert invalid_label in error.message


def test_label_validation_error_keeps_base_message_for_non_identity_cases(
    monkeypatch: pytest.MonkeyPatch,
):
    _set_caller_identity(monkeypatch, "user_email")
    non_target_metric = LabelValidationError(
        metric_name="litellm_overhead_with_guardrails_latency_metric",
        invalid_labels=["api_key_alias"],
        valid_labels=[],
    )
    non_identity_label = LabelValidationError(
        metric_name="litellm_deployment_total_requests",
        invalid_labels=["bogus_label"],
        valid_labels=[],
    )

    for error in (non_target_metric, non_identity_label):
        assert "caller-identity" not in error.message
        assert error.message.startswith("Invalid labels for metric")