mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
581 lines
21 KiB
Python
581 lines
21 KiB
Python
"""
|
|
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
|
|
|
|
def test_record_guardrail_fail_open_increments_errors_total(self):
|
|
prometheus_logger = PrometheusLogger()
|
|
|
|
prometheus_logger.record_guardrail_fail_open(guardrail_name="a365", hook_type="pre_call")
|
|
prometheus_logger.record_guardrail_fail_open(guardrail_name="a365", hook_type="pre_call")
|
|
|
|
assert (
|
|
REGISTRY.get_sample_value(
|
|
"litellm_guardrail_errors_total",
|
|
{"guardrail_name": "a365", "error_type": "fail_open", "hook_type": "pre_call"},
|
|
)
|
|
== 2.0
|
|
)
|
|
assert (
|
|
REGISTRY.get_sample_value(
|
|
"litellm_guardrail_requests_total",
|
|
{"guardrail_name": "a365", "status": "success", "hook_type": "pre_call"},
|
|
)
|
|
is None
|
|
)
|
|
|
|
def test_record_guardrail_fail_open_swallows_metric_errors(self):
|
|
prometheus_logger = PrometheusLogger()
|
|
broken_counter = MagicMock()
|
|
broken_counter.labels.side_effect = ValueError("registry exploded")
|
|
prometheus_logger.litellm_guardrail_errors_total = broken_counter
|
|
|
|
prometheus_logger.record_guardrail_fail_open(guardrail_name="a365", hook_type="pre_call")
|
|
|
|
broken_counter.labels.assert_called_once_with(
|
|
guardrail_name="a365", error_type="fail_open", hook_type="pre_call"
|
|
)
|