diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index aa06ec99eb7..9aea69c34a6 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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, diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index cdcc8807617..6a696345f9c 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -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 diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index 119fe3cd46a..254ab9f5a5c 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -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() -