scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

test: fix tests

Browse source
This commit is contained in:
Krrish Dholakia 2025-06-19 16:35:05 -07:00
parent 677ef3cdde
commit 3540984965
3 changed files with 76 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/integrations/prometheus.py
View file

@ -551,6 +551,7 @@ class PrometheusLogger(CustomLogger):
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
requested_model=standard_logging_payload["model_group"],
model_group=standard_logging_payload["model_group"],
team=user_api_team,
team_alias=user_api_team_alias,
user=user_id,

3
litellm/types/integrations/prometheus.py
View file

@ -336,6 +336,9 @@ class UserAPIKeyLabelValues(BaseModel):
team_alias: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.TEAM_ALIAS.value)
] = None
model_group: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.MODEL_GROUP.value)
] = None
requested_model: Annotated[
Optional[str], Field(..., alias=UserAPIKeyLabelNames.REQUESTED_MODEL.value)
] = None

118
tests/logging_callback_tests/test_prometheus_unit_tests.py
View file

@ -219,7 +219,14 @@ def test_increment_token_metrics(prometheus_logger):
)
prometheus_logger.litellm_tokens_metric.labels.assert_called_once_with(
end_user=None, user=None, hashed_api_key='test_hash', api_key_alias='test_alias', team='test_team', team_alias='test_team_alias', requested_model=None, model='gpt-3.5-turbo'
end_user=None,
user=None,
hashed_api_key="test_hash",
api_key_alias="test_alias",
team="test_team",
team_alias="test_team_alias",
requested_model=None,
model="gpt-3.5-turbo",
)
prometheus_logger.litellm_tokens_metric.labels().inc.assert_called_once_with(100)
@ -836,12 +843,12 @@ def test_set_llm_deployment_success_metrics(prometheus_logger):
# Verify remaining requests metric
prometheus_logger.litellm_remaining_requests_metric.labels.assert_called_once_with(
api_base="https://api.openai.com",
api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"],
api_provider="openai",
model_group="my_custom_model_group", # model_group / requested model from create_standard_logging_payload()
api_provider="openai", # llm provider
api_base="https://api.openai.com", # api base
litellm_model_name="gpt-3.5-turbo", # actual model used - litellm model name
hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"],
litellm_model_name="gpt-3.5-turbo",
requested_model="my_custom_model_group",
api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"],
)
prometheus_logger.litellm_remaining_requests_metric.labels().set.assert_called_once_with(
@ -855,7 +862,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger):
api_provider="openai",
hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"],
litellm_model_name="gpt-3.5-turbo",
requested_model="my_custom_model_group",
model_group="my_custom_model_group",
)
prometheus_logger.litellm_remaining_tokens_metric.labels().set.assert_called_once_with(
@ -915,7 +922,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger):
api_provider="openai",
hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"],
litellm_model_name="gpt-3.5-turbo",
requested_model="my_custom_model_group",
model_group="my_custom_model_group",
)
# Calculate expected latency per token (1 second / 10 tokens = 0.1 seconds per token)
@ -1499,60 +1506,66 @@ def test_get_exception_class_name(prometheus_logger):
"""
# Test case 1: Exception with llm_provider
rate_limit_error = litellm.RateLimitError(
message="Rate limit exceeded",
llm_provider="openai",
model="gpt-3.5-turbo"
message="Rate limit exceeded", llm_provider="openai", model="gpt-3.5-turbo"
)
assert (
prometheus_logger._get_exception_class_name(rate_limit_error)
== "Openai.RateLimitError"
)
assert prometheus_logger._get_exception_class_name(rate_limit_error) == "Openai.RateLimitError"
# Test case 2: Exception with empty llm_provider
auth_error = litellm.AuthenticationError(
message="Invalid API key",
llm_provider="",
model="gpt-4"
message="Invalid API key", llm_provider="", model="gpt-4"
)
assert (
prometheus_logger._get_exception_class_name(auth_error) == "AuthenticationError"
)
assert prometheus_logger._get_exception_class_name(auth_error) == "AuthenticationError"
# Test case 3: Exception with None llm_provider
context_window_error = litellm.ContextWindowExceededError(
message="Context length exceeded",
llm_provider=None,
model="gpt-4"
message="Context length exceeded", llm_provider=None, model="gpt-4"
)
assert (
prometheus_logger._get_exception_class_name(context_window_error)
== "ContextWindowExceededError"
)
assert prometheus_logger._get_exception_class_name(context_window_error) == "ContextWindowExceededError"
def test_set_llm_deployment_success_metrics_with_label_filtering():
"""
Test that set_llm_deployment_success_metrics correctly uses prometheus_label_factory
Test that set_llm_deployment_success_metrics correctly uses prometheus_label_factory
and respects label filtering configuration to prevent "Incorrect label names" errors.
"""
from litellm.types.integrations.prometheus import PrometheusMetricsConfig
# Create a prometheus logger with label filtering configuration
config = [
PrometheusMetricsConfig(
group="test_group",
metrics=[
"litellm_overhead_latency_metric",
"litellm_remaining_requests_metric",
"litellm_remaining_requests_metric",
"litellm_remaining_tokens_metric",
"litellm_deployment_success_responses",
"litellm_deployment_total_requests"
"litellm_deployment_total_requests",
],
include_labels=["requested_model", "api_provider", "hashed_api_key"] # Limited labels
include_labels=[
"requested_model",
"api_provider",
"hashed_api_key",
], # Limited labels
)
]
# Mock litellm.prometheus_metrics_config
with patch('litellm.prometheus_metrics_config', config):
with patch("litellm.prometheus_metrics_config", config):
# Clear registry before creating new logger
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
prometheus_logger = PrometheusLogger()
# Mock all the metrics used in the method
prometheus_logger.litellm_overhead_latency_metric = MagicMock()
prometheus_logger.litellm_remaining_requests_metric = MagicMock()
@ -1607,46 +1620,59 @@ def test_set_llm_deployment_success_metrics_with_label_filtering():
# Verify that metrics were called with filtered labels (only the configured ones)
# The exact labels depend on what get_labels_for_metric returns for each metric
# Verify overhead latency metric was called with filtered labels
prometheus_logger.litellm_overhead_latency_metric.labels.assert_called_once()
overhead_labels = prometheus_logger.litellm_overhead_latency_metric.labels.call_args[1]
overhead_labels = (
prometheus_logger.litellm_overhead_latency_metric.labels.call_args[1]
)
# Should only contain the filtered labels that are supported for this metric
expected_filtered_labels = {"requested_model", "api_provider", "hashed_api_key"}
actual_labels = set(k for k in overhead_labels.keys() if k is not None)
# Verify that only expected labels are present (subset of configured labels)
assert actual_labels <= expected_filtered_labels
# Verify remaining requests metric was called with filtered labels
# Verify remaining requests metric was called with filtered labels
prometheus_logger.litellm_remaining_requests_metric.labels.assert_called_once()
requests_labels = prometheus_logger.litellm_remaining_requests_metric.labels.call_args[1]
requests_labels = (
prometheus_logger.litellm_remaining_requests_metric.labels.call_args[1]
)
actual_labels = set(k for k in requests_labels.keys() if k is not None)
assert actual_labels <= expected_filtered_labels
# Verify remaining tokens metric was called with filtered labels
prometheus_logger.litellm_remaining_tokens_metric.labels.assert_called_once()
tokens_labels = prometheus_logger.litellm_remaining_tokens_metric.labels.call_args[1]
tokens_labels = (
prometheus_logger.litellm_remaining_tokens_metric.labels.call_args[1]
)
actual_labels = set(k for k in tokens_labels.keys() if k is not None)
assert actual_labels <= expected_filtered_labels
# Verify deployment success responses metric was called with filtered labels
prometheus_logger.litellm_deployment_success_responses.labels.assert_called_once()
success_labels = prometheus_logger.litellm_deployment_success_responses.labels.call_args[1]
success_labels = (
prometheus_logger.litellm_deployment_success_responses.labels.call_args[1]
)
actual_labels = set(k for k in success_labels.keys() if k is not None)
assert actual_labels <= expected_filtered_labels
# Verify deployment total requests metric was called with filtered labels
prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once()
total_labels = prometheus_logger.litellm_deployment_total_requests.labels.call_args[1]
total_labels = (
prometheus_logger.litellm_deployment_total_requests.labels.call_args[1]
)
actual_labels = set(total_labels.keys())
assert actual_labels.issubset(expected_filtered_labels.union({None}))
# Verify all metrics were actually called (no exceptions were raised)
prometheus_logger.litellm_overhead_latency_metric.labels().observe.assert_called_once()
prometheus_logger.litellm_remaining_requests_metric.labels().set.assert_called_once_with(123)
prometheus_logger.litellm_remaining_tokens_metric.labels().set.assert_called_once_with(4321)
prometheus_logger.litellm_remaining_requests_metric.labels().set.assert_called_once_with(
123
)
prometheus_logger.litellm_remaining_tokens_metric.labels().set.assert_called_once_with(
4321
)
prometheus_logger.litellm_deployment_success_responses.labels().inc.assert_called_once()
prometheus_logger.litellm_deployment_total_requests.labels().inc.assert_called_once()