""" Unit tests for prometheus queue time and guardrail metrics """ from datetime import datetime from unittest.mock import MagicMock import pytest from prometheus_client import REGISTRY from litellm.integrations.prometheus import PrometheusLogger from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @pytest.fixture(autouse=True) def cleanup_prometheus_registry(): """Clean up prometheus registry between tests""" # Clear the registry before each test collectors = list(REGISTRY._collector_to_names.keys()) for collector in collectors: REGISTRY.unregister(collector) yield # Clean up after test collectors = list(REGISTRY._collector_to_names.keys()) for collector in collectors: REGISTRY.unregister(collector) class TestPrometheusQueueTimeMetric: """Test request queue time metric recording""" def test_queue_time_metric_recorded_in_set_latency_metrics(self): """Test that queue time metric is recorded when queue_time_seconds is present in metadata""" # Arrange prometheus_logger = PrometheusLogger() # Mock the metric mock_metric = MagicMock() mock_labeled_metric = MagicMock() mock_metric.labels.return_value = mock_labeled_metric prometheus_logger.litellm_request_queue_time_metric = mock_metric # Create mock kwargs with queue_time_seconds in metadata queue_time_seconds = 0.5 kwargs = { "litellm_params": {"metadata": {"queue_time_seconds": queue_time_seconds}}, "model": "gpt-3.5-turbo", "start_time": datetime.now(), "end_time": datetime.now(), } enum_values = UserAPIKeyLabelValues( end_user=None, hashed_api_key="test-key", api_key_alias="test-alias", requested_model="gpt-3.5-turbo", model_group="gpt-3.5-turbo", team=None, team_alias=None, user=None, user_email=None, status_code="200", model="gpt-3.5-turbo", litellm_model_name="gpt-3.5-turbo", tags=[], model_id="gpt-3.5-turbo", api_base="https://api.openai.com", api_provider="openai", exception_status=None, exception_class=None, custom_metadata_labels={}, route=None, ) # Act prometheus_logger._set_latency_metrics( kwargs=kwargs, model="gpt-3.5-turbo", user_api_key="test-key", user_api_key_alias="test-alias", user_api_team=None, user_api_team_alias=None, enum_values=enum_values, ) # Assert - queue time metric should be called mock_metric.labels.assert_called() # Check that observe was called on the queue time metric assert mock_labeled_metric.observe.called # Verify the observed value observed_value = None for call in mock_labeled_metric.observe.call_args_list: if len(call[0]) > 0: observed_value = call[0][0] if observed_value == queue_time_seconds: break assert observed_value == queue_time_seconds assert observed_value >= 0 def test_queue_time_metric_not_recorded_when_missing(self): """Test that queue time metric is not recorded when queue_time_seconds is missing""" # Arrange prometheus_logger = PrometheusLogger() # Mock the metric mock_metric = MagicMock() mock_labeled_metric = MagicMock() mock_metric.labels.return_value = mock_labeled_metric prometheus_logger.litellm_request_queue_time_metric = mock_metric # Create mock kwargs without queue_time_seconds kwargs = { "litellm_params": {"metadata": {}}, "model": "gpt-3.5-turbo", "start_time": datetime.now(), "end_time": datetime.now(), } enum_values = UserAPIKeyLabelValues( end_user=None, hashed_api_key="test-key", api_key_alias="test-alias", requested_model="gpt-3.5-turbo", model_group="gpt-3.5-turbo", team=None, team_alias=None, user=None, user_email=None, status_code="200", model="gpt-3.5-turbo", litellm_model_name="gpt-3.5-turbo", tags=[], model_id="gpt-3.5-turbo", api_base="https://api.openai.com", api_provider="openai", exception_status=None, exception_class=None, custom_metadata_labels={}, route=None, ) # Act prometheus_logger._set_latency_metrics( kwargs=kwargs, model="gpt-3.5-turbo", user_api_key="test-key", user_api_key_alias="test-alias", user_api_team=None, user_api_team_alias=None, enum_values=enum_values, ) # Assert - queue time metric should not be called (queue_time_seconds is None) # We check that observe was not called with queue_time_seconds queue_time_called = False for call in mock_labeled_metric.observe.call_args_list: if len(call[0]) > 0 and call[0][0] == 0.5: # Our test queue time value queue_time_called = True break assert ( not queue_time_called ), "Queue time metric should not be recorded when queue_time_seconds is missing" def test_queue_time_metric_not_recorded_when_negative(self): """Test that queue time metric is not recorded when queue_time_seconds is negative""" # Arrange prometheus_logger = PrometheusLogger() # Mock the metric mock_metric = MagicMock() mock_labeled_metric = MagicMock() mock_metric.labels.return_value = mock_labeled_metric prometheus_logger.litellm_request_queue_time_metric = mock_metric # Create mock kwargs with negative queue_time_seconds kwargs = { "litellm_params": { "metadata": {"queue_time_seconds": -0.1} # Negative value }, "model": "gpt-3.5-turbo", "start_time": datetime.now(), "end_time": datetime.now(), } enum_values = UserAPIKeyLabelValues( end_user=None, hashed_api_key="test-key", api_key_alias="test-alias", requested_model="gpt-3.5-turbo", model_group="gpt-3.5-turbo", team=None, team_alias=None, user=None, user_email=None, status_code="200", model="gpt-3.5-turbo", litellm_model_name="gpt-3.5-turbo", tags=[], model_id="gpt-3.5-turbo", api_base="https://api.openai.com", api_provider="openai", exception_status=None, exception_class=None, custom_metadata_labels={}, route=None, ) # Act prometheus_logger._set_latency_metrics( kwargs=kwargs, model="gpt-3.5-turbo", user_api_key="test-key", user_api_key_alias="test-alias", user_api_team=None, user_api_team_alias=None, enum_values=enum_values, ) # Assert - queue time metric should not be called for negative values # We check that observe was not called with the negative value negative_value_called = False for call in mock_labeled_metric.observe.call_args_list: if len(call[0]) > 0 and call[0][0] == -0.1: negative_value_called = True break assert ( not negative_value_called ), "Queue time metric should not be recorded for negative values" class TestPrometheusTotalLatencyMetric: """litellm_request_total_latency_metric must be true end-to-end latency: start_time (set after auth already completed, see LIT-6012) plus queue_time_seconds (the auth + pre-call setup window queue_time_seconds itself covers), not start_time alone.""" @staticmethod def _enum_values() -> UserAPIKeyLabelValues: return UserAPIKeyLabelValues( end_user=None, hashed_api_key="test-key", api_key_alias="test-alias", requested_model="gpt-3.5-turbo", model_group="gpt-3.5-turbo", team=None, team_alias=None, user=None, user_email=None, status_code="200", model="gpt-3.5-turbo", litellm_model_name="gpt-3.5-turbo", tags=[], model_id="gpt-3.5-turbo", api_base="https://api.openai.com", api_provider="openai", exception_status=None, exception_class=None, custom_metadata_labels={}, route=None, ) def test_total_latency_includes_queue_time_when_present(self): """The observed total-latency value must be (end_time - start_time) + queue_time_seconds, so auth/pre-call time (queue_time_seconds) is not silently excluded from "total" latency.""" prometheus_logger = PrometheusLogger() mock_metric = MagicMock() mock_labeled_metric = MagicMock() mock_metric.labels.return_value = mock_labeled_metric prometheus_logger.litellm_request_total_latency_metric = mock_metric start_time = datetime(2024, 1, 1, 0, 0, 0) end_time = datetime(2024, 1, 1, 0, 0, 2) # 2.0s of LLM-call/post-call time queue_time_seconds = 0.5 # auth + pre-call setup time kwargs = { "litellm_params": {"metadata": {"queue_time_seconds": queue_time_seconds}}, "model": "gpt-3.5-turbo", "start_time": start_time, "end_time": end_time, } prometheus_logger._set_latency_metrics( kwargs=kwargs, model="gpt-3.5-turbo", user_api_key="test-key", user_api_key_alias="test-alias", user_api_team=None, user_api_team_alias=None, enum_values=self._enum_values(), ) observed_value = mock_labeled_metric.observe.call_args_list[0][0][0] assert observed_value == pytest.approx(2.5) def test_total_latency_falls_back_to_start_end_delta_without_queue_time(self): """Without queue_time_seconds (e.g. a non-proxy caller), the metric must still observe the plain end_time - start_time delta rather than erroring or dropping it.""" prometheus_logger = PrometheusLogger() mock_metric = MagicMock() mock_labeled_metric = MagicMock() mock_metric.labels.return_value = mock_labeled_metric prometheus_logger.litellm_request_total_latency_metric = mock_metric start_time = datetime(2024, 1, 1, 0, 0, 0) end_time = datetime(2024, 1, 1, 0, 0, 2) kwargs = { "litellm_params": {"metadata": {}}, "model": "gpt-3.5-turbo", "start_time": start_time, "end_time": end_time, } prometheus_logger._set_latency_metrics( kwargs=kwargs, model="gpt-3.5-turbo", user_api_key="test-key", user_api_key_alias="test-alias", user_api_team=None, user_api_team_alias=None, enum_values=self._enum_values(), ) observed_value = mock_labeled_metric.observe.call_args_list[0][0][0] assert observed_value == pytest.approx(2.0) def test_total_latency_ignores_negative_queue_time(self): """A negative queue_time_seconds (clock skew / bad data) must not be added in -- matches the existing >= 0 guard on the standalone queue-time metric.""" prometheus_logger = PrometheusLogger() mock_metric = MagicMock() mock_labeled_metric = MagicMock() mock_metric.labels.return_value = mock_labeled_metric prometheus_logger.litellm_request_total_latency_metric = mock_metric start_time = datetime(2024, 1, 1, 0, 0, 0) end_time = datetime(2024, 1, 1, 0, 0, 2) kwargs = { "litellm_params": {"metadata": {"queue_time_seconds": -0.1}}, "model": "gpt-3.5-turbo", "start_time": start_time, "end_time": end_time, } prometheus_logger._set_latency_metrics( kwargs=kwargs, model="gpt-3.5-turbo", user_api_key="test-key", user_api_key_alias="test-alias", user_api_team=None, user_api_team_alias=None, enum_values=self._enum_values(), ) observed_value = mock_labeled_metric.observe.call_args_list[0][0][0] assert observed_value == pytest.approx(2.0) class TestPrometheusGuardrailMetrics: """Test guardrail metrics recording""" def test_record_guardrail_metrics_success(self): """Test recording guardrail metrics for successful execution""" # Arrange prometheus_logger = PrometheusLogger() # Mock metrics mock_latency_metric = MagicMock() mock_requests_metric = MagicMock() mock_errors_metric = MagicMock() prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric guardrail_name = "test_guardrail" latency_seconds = 0.15 status = "success" error_type = None hook_type = "pre_call" # Act prometheus_logger._record_guardrail_metrics( guardrail_name=guardrail_name, latency_seconds=latency_seconds, status=status, error_type=error_type, hook_type=hook_type, ) # Assert - latency metric should be recorded mock_latency_metric.labels.assert_called_once_with( guardrail_name=guardrail_name, status=status, error_type="none", hook_type=hook_type, ) mock_latency_metric.labels.return_value.observe.assert_called_once_with( latency_seconds ) # Assert - requests metric should be incremented mock_requests_metric.labels.assert_called_once_with( guardrail_name=guardrail_name, status=status, hook_type=hook_type, ) mock_requests_metric.labels.return_value.inc.assert_called_once() # Assert - errors metric should NOT be called for success mock_errors_metric.labels.assert_not_called() def test_record_guardrail_metrics_error(self): """Test recording guardrail metrics for failed execution""" # Arrange prometheus_logger = PrometheusLogger() # Mock metrics mock_latency_metric = MagicMock() mock_requests_metric = MagicMock() mock_errors_metric = MagicMock() prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric prometheus_logger.litellm_guardrail_errors_total = mock_errors_metric guardrail_name = "test_guardrail" latency_seconds = 0.2 status = "error" error_type = "ValueError" hook_type = "pre_call" # Act prometheus_logger._record_guardrail_metrics( guardrail_name=guardrail_name, latency_seconds=latency_seconds, status=status, error_type=error_type, hook_type=hook_type, ) # Assert - latency metric should be recorded mock_latency_metric.labels.assert_called_once_with( guardrail_name=guardrail_name, status=status, error_type=error_type, hook_type=hook_type, ) mock_latency_metric.labels.return_value.observe.assert_called_once_with( latency_seconds ) # Assert - requests metric should be incremented mock_requests_metric.labels.assert_called_once_with( guardrail_name=guardrail_name, status=status, hook_type=hook_type, ) mock_requests_metric.labels.return_value.inc.assert_called_once() # Assert - errors metric should be incremented mock_errors_metric.labels.assert_called_once_with( guardrail_name=guardrail_name, error_type=error_type, hook_type=hook_type, ) mock_errors_metric.labels.return_value.inc.assert_called_once() def test_record_guardrail_metrics_during_call_hook(self): """Test recording guardrail metrics for during_call hook""" # Arrange prometheus_logger = PrometheusLogger() # Mock metrics mock_latency_metric = MagicMock() mock_requests_metric = MagicMock() prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric guardrail_name = "moderation_guardrail" latency_seconds = 0.1 status = "success" hook_type = "during_call" # Act prometheus_logger._record_guardrail_metrics( guardrail_name=guardrail_name, latency_seconds=latency_seconds, status=status, error_type=None, hook_type=hook_type, ) # Assert - hook_type should be "during_call" mock_latency_metric.labels.assert_called_once() call_kwargs = mock_latency_metric.labels.call_args[1] assert call_kwargs["hook_type"] == "during_call" def test_record_guardrail_metrics_handles_exception(self): """Test that _record_guardrail_metrics handles exceptions gracefully""" # Arrange prometheus_logger = PrometheusLogger() # Mock metric to raise exception mock_metric = MagicMock() mock_metric.labels.side_effect = Exception("Test error") prometheus_logger.litellm_guardrail_latency_metric = mock_metric prometheus_logger.litellm_guardrail_requests_total = MagicMock() # Act & Assert - should not raise exception try: prometheus_logger._record_guardrail_metrics( guardrail_name="test", latency_seconds=0.1, status="success", error_type=None, hook_type="pre_call", ) except Exception: pytest.fail("_record_guardrail_metrics should handle exceptions gracefully") def test_record_guardrail_metrics_with_guardrail_name_attribute(self): """Test that guardrail name is extracted from guardrail_name attribute if available""" # Arrange prometheus_logger = PrometheusLogger() # Mock metrics mock_latency_metric = MagicMock() mock_requests_metric = MagicMock() prometheus_logger.litellm_guardrail_latency_metric = mock_latency_metric prometheus_logger.litellm_guardrail_requests_total = mock_requests_metric guardrail_name = "custom_guardrail_name" latency_seconds = 0.1 status = "success" hook_type = "pre_call" # Act prometheus_logger._record_guardrail_metrics( guardrail_name=guardrail_name, latency_seconds=latency_seconds, status=status, error_type=None, hook_type=hook_type, ) # Assert - guardrail_name should be used mock_latency_metric.labels.assert_called_once() call_kwargs = mock_latency_metric.labels.call_args[1] assert call_kwargs["guardrail_name"] == guardrail_name