diff --git a/tests/test_litellm/integrations/test_prometheus_latency_buckets.py b/tests/test_litellm/integrations/test_prometheus_latency_buckets.py new file mode 100644 index 00000000000..3160bb5a9de --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_latency_buckets.py @@ -0,0 +1,144 @@ +""" +Tests that prometheus_latency_buckets is respected end-to-end: +- histograms are registered with the custom boundaries +- observed values land in the expected bucket +""" + +from datetime import datetime, timedelta + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +@pytest.fixture(autouse=True) +def reset_prometheus_registry(): + collectors = list(REGISTRY._collector_to_names.keys()) + for c in collectors: + REGISTRY.unregister(c) + yield + collectors = list(REGISTRY._collector_to_names.keys()) + for c in collectors: + REGISTRY.unregister(c) + + +@pytest.fixture(autouse=True) +def reset_litellm_buckets(): + original = litellm.prometheus_latency_buckets + yield + litellm.prometheus_latency_buckets = original + + +def _make_enum_values() -> UserAPIKeyLabelValues: + return UserAPIKeyLabelValues( + end_user=None, + hashed_api_key="test-key", + api_key_alias="test-alias", + requested_model="gpt-4", + model_group="gpt-4", + team=None, + team_alias=None, + user=None, + user_email=None, + status_code="200", + model="gpt-4", + litellm_model_name="gpt-4", + tags=[], + model_id="model-id-1", + api_base="https://api.openai.com", + api_provider="openai", + exception_status=None, + exception_class=None, + custom_metadata_labels={}, + route=None, + ) + + +def _get_bucket_count(metric_name: str, le: str) -> float: + """Return the bucket count for a specific upper bound from the live registry.""" + for metric_family in REGISTRY.collect(): + if metric_family.name == metric_name: + for sample in metric_family.samples: + if ( + sample.name == f"{metric_name}_bucket" + and sample.labels.get("le") == le + ): + return sample.value + raise AssertionError(f"Bucket le={le} not found in metric {metric_name}") + + +def test_custom_latency_buckets_registered(): + """Histograms use prometheus_latency_buckets when set.""" + litellm.prometheus_latency_buckets = (1.0, 5.0, float("inf")) + + logger = PrometheusLogger() + + assert logger._get_latency_buckets() == (1.0, 5.0, float("inf")) + assert logger.litellm_request_total_latency_metric._upper_bounds == [ + 1.0, + 5.0, + float("inf"), + ] + assert logger.litellm_llm_api_latency_metric._upper_bounds == [ + 1.0, + 5.0, + float("inf"), + ] + + +def test_observation_lands_in_correct_custom_bucket(): + """ + With buckets [1.0, 5.0, inf], a 2-second observation must: + - NOT be counted in the le=1.0 bucket + - BE counted in the le=5.0 bucket + - BE counted in the le=+Inf bucket + """ + litellm.prometheus_latency_buckets = (1.0, 5.0, float("inf")) + + logger = PrometheusLogger() + enum_values = _make_enum_values() + + now = datetime.now() + kwargs = { + "start_time": now - timedelta(seconds=2), + "end_time": now, + "api_call_start_time": now - timedelta(seconds=2), + "litellm_params": {"metadata": {}}, + "model": "gpt-4", + } + + logger._set_latency_metrics( + kwargs=kwargs, + model="gpt-4", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=enum_values, + ) + + # le=1.0 bucket must be 0 (2s > 1.0) + assert _get_bucket_count("litellm_request_total_latency_metric", "1.0") == 0.0 + + # le=5.0 bucket must be 1 (2s <= 5.0) + assert _get_bucket_count("litellm_request_total_latency_metric", "5.0") == 1.0 + + # le=+Inf must also be 1 + assert _get_bucket_count("litellm_request_total_latency_metric", "+Inf") == 1.0 + + +def test_default_buckets_used_when_not_set(): + """When prometheus_latency_buckets is None, the module-level LATENCY_BUCKETS constant is used.""" + from litellm.types.integrations.prometheus import LATENCY_BUCKETS + + litellm.prometheus_latency_buckets = None + + logger = PrometheusLogger() + + assert logger._get_latency_buckets() == LATENCY_BUCKETS + assert logger.litellm_request_total_latency_metric._upper_bounds == list( + LATENCY_BUCKETS + )