diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index c01f7481277..2ec2f41b3bb 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -14,6 +14,7 @@ from typing import ( Literal, Optional, Tuple, + Union, cast, ) @@ -44,6 +45,7 @@ def _get_cached_end_user_id_for_cost_tracking(): global _get_end_user_id_for_cost_tracking if _get_end_user_id_for_cost_tracking is None: from litellm.utils import get_end_user_id_for_cost_tracking + _get_end_user_id_for_cost_tracking = get_end_user_id_for_cost_tracking return _get_end_user_id_for_cost_tracking @@ -329,6 +331,25 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_requests_metric"), ) + # Cache metrics + self.litellm_cache_hits_metric = self._counter_factory( + name="litellm_cache_hits_metric", + documentation="Total number of LiteLLM cache hits", + labelnames=self.get_labels_for_metric("litellm_cache_hits_metric"), + ) + + self.litellm_cache_misses_metric = self._counter_factory( + name="litellm_cache_misses_metric", + documentation="Total number of LiteLLM cache misses", + labelnames=self.get_labels_for_metric("litellm_cache_misses_metric"), + ) + + self.litellm_cached_tokens_metric = self._counter_factory( + name="litellm_cached_tokens_metric", + documentation="Total tokens served from LiteLLM cache", + labelnames=self.get_labels_for_metric("litellm_cached_tokens_metric"), + ) + except Exception as e: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e @@ -795,7 +816,7 @@ class PrometheusLogger(CustomLogger): litellm_params = kwargs.get("litellm_params", {}) or {} _metadata = litellm_params.get("metadata", {}) get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - + end_user_id = get_end_user_id_for_cost_tracking( litellm_params, service_type="prometheus" ) @@ -815,7 +836,7 @@ class PrometheusLogger(CustomLogger): user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[ "metadata" ].get("user_api_key_auth_metadata") - + # Include top-level metadata fields (excluding nested dictionaries) # This allows accessing fields like requester_ip_address from top-level metadata top_level_metadata = standard_logging_payload.get("metadata", {}) @@ -826,7 +847,7 @@ class PrometheusLogger(CustomLogger): for k, v in top_level_metadata.items() if not isinstance(v, dict) # Exclude nested dicts to avoid conflicts } - + combined_metadata: Dict[str, Any] = { **top_level_fields, # Include top-level fields first **(_requester_metadata if _requester_metadata else {}), @@ -945,6 +966,12 @@ class PrometheusLogger(CustomLogger): kwargs, start_time, end_time, enum_values, output_tokens ) + # cache metrics + self._increment_cache_metrics( + standard_logging_payload=standard_logging_payload, # type: ignore + enum_values=enum_values, + ) + if ( standard_logging_payload["stream"] is True ): # log successful streaming requests from logging event hook. @@ -1014,6 +1041,54 @@ class PrometheusLogger(CustomLogger): standard_logging_payload["completion_tokens"] ) + def _increment_cache_metrics( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + ): + """ + Increment cache-related Prometheus metrics based on cache hit/miss status. + + Args: + standard_logging_payload: Contains cache_hit field (True/False/None) + enum_values: Label values for Prometheus metrics + """ + cache_hit = standard_logging_payload.get("cache_hit") + + # Only track if cache_hit has a definite value (True or False) + if cache_hit is None: + return + + if cache_hit is True: + # Increment cache hits counter + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_cache_hits_metric" + ), + enum_values=enum_values, + ) + self.litellm_cache_hits_metric.labels(**_labels).inc() + + # Increment cached tokens counter + total_tokens = standard_logging_payload.get("total_tokens", 0) + if total_tokens > 0: + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_cached_tokens_metric" + ), + enum_values=enum_values, + ) + self.litellm_cached_tokens_metric.labels(**_labels).inc(total_tokens) + else: + # cache_hit is False - increment cache misses counter + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_cache_misses_metric" + ), + enum_values=enum_values, + ) + self.litellm_cache_misses_metric.labels(**_labels).inc() + async def _increment_remaining_budget_metrics( self, user_api_team: Optional[str], @@ -1196,7 +1271,7 @@ class PrometheusLogger(CustomLogger): ) litellm_params = kwargs.get("litellm_params", {}) or {} get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - + end_user_id = get_end_user_id_for_cost_tracking( litellm_params, service_type="prometheus" ) @@ -1398,7 +1473,6 @@ class PrometheusLogger(CustomLogger): api_provider=llm_provider or "", ) if exception is not None: - _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( metric_name="litellm_deployment_failure_responses" @@ -1431,12 +1505,11 @@ class PrometheusLogger(CustomLogger): enum_values: UserAPIKeyLabelValues, output_tokens: float = 1.0, ): - try: verbose_logger.debug("setting remaining tokens requests metric") - standard_logging_payload: Optional[StandardLoggingPayload] = ( - request_kwargs.get("standard_logging_object") - ) + standard_logging_payload: Optional[ + StandardLoggingPayload + ] = request_kwargs.get("standard_logging_object") if standard_logging_payload is None: return @@ -2208,10 +2281,10 @@ class PrometheusLogger(CustomLogger): from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - prometheus_loggers: List[CustomLogger] = ( - litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=PrometheusLogger - ) + prometheus_loggers: List[ + CustomLogger + ] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=PrometheusLogger ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers)) @@ -2283,7 +2356,7 @@ def prometheus_label_factory( if UserAPIKeyLabelNames.END_USER.value in filtered_labels: get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking() - + filtered_labels["end_user"] = get_end_user_id_for_cost_tracking( litellm_params={"user_api_key_end_user_id": enum_values.end_user}, service_type="prometheus", diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 6a254fc8252..fb439a9541b 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,7 +1,7 @@ import re from dataclasses import dataclass from enum import Enum -from typing import Dict, List, Literal, Optional, Tuple, Union +from typing import Dict, List, Literal, Optional, Tuple from pydantic import BaseModel, Field from typing_extensions import Annotated @@ -185,6 +185,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_redis_daily_spend_update_queue_size", "litellm_in_memory_spend_update_queue_size", "litellm_redis_spend_update_queue_size", + # Cache metrics + "litellm_cache_hits_metric", + "litellm_cache_misses_metric", + "litellm_cached_tokens_metric", ] @@ -436,6 +440,21 @@ class PrometheusMetricLabels: litellm_redis_spend_update_queue_size: List[str] = [] + # Cache metrics - track cache hits, misses, and tokens served from cache + _cache_metric_labels = [ + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.END_USER.value, + UserAPIKeyLabelNames.USER.value, + ] + + litellm_cache_hits_metric = _cache_metric_labels + litellm_cache_misses_metric = _cache_metric_labels + litellm_cached_tokens_metric = _cache_metric_labels + @staticmethod def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]: default_labels = getattr(PrometheusMetricLabels, label_name) diff --git a/tests/test_litellm/integrations/test_prometheus_cache_metrics.py b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py new file mode 100644 index 00000000000..660757673f6 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_cache_metrics.py @@ -0,0 +1,211 @@ +""" +Unit tests for cache Prometheus metrics. + +Run with: poetry run pytest tests/test_litellm/integrations/test_prometheus_cache_metrics.py -v +""" +import pytest +from unittest.mock import MagicMock, patch +from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + +class TestPrometheusCacheMetrics: + """Tests for cache-related Prometheus metrics""" + + @pytest.fixture + def sample_enum_values(self): + """Create sample enum values for labels""" + return UserAPIKeyLabelValues( + end_user="test-end-user", + hashed_api_key="test-key-hash", + api_key_alias="test-key-alias", + team="test-team", + team_alias="test-team-alias", + user="test-user", + model="gpt-3.5-turbo", + ) + + def test_cache_metrics_defined_in_types(self): + """Test that cache metrics are defined in DEFINED_PROMETHEUS_METRICS""" + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + from typing import get_args + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + assert "litellm_cache_hits_metric" in defined_metrics + assert "litellm_cache_misses_metric" in defined_metrics + assert "litellm_cached_tokens_metric" in defined_metrics + + def test_cache_metric_labels_defined(self): + """Test that cache metric labels are properly defined""" + from litellm.types.integrations.prometheus import PrometheusMetricLabels + + # Verify labels are defined for each cache metric + assert hasattr(PrometheusMetricLabels, "litellm_cache_hits_metric") + assert hasattr(PrometheusMetricLabels, "litellm_cache_misses_metric") + assert hasattr(PrometheusMetricLabels, "litellm_cached_tokens_metric") + + # Verify labels include expected keys + expected_labels = [ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + for label in expected_labels: + assert label in PrometheusMetricLabels.litellm_cache_hits_metric + assert label in PrometheusMetricLabels.litellm_cache_misses_metric + assert label in PrometheusMetricLabels.litellm_cached_tokens_metric + + def test_increment_cache_metrics_on_cache_hit(self, sample_enum_values): + """Test that cache hit increments the correct metrics""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + # Import the method directly and bind it to our mock + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=True + standard_logging_payload = { + "cache_hit": True, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method using unbound method approach + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify cache hits metric was incremented + mock_logger.litellm_cache_hits_metric.labels.assert_called() + mock_logger.litellm_cache_hits_metric.labels().inc.assert_called_once() + + # Verify cached tokens metric was incremented with total_tokens + mock_logger.litellm_cached_tokens_metric.labels.assert_called() + mock_logger.litellm_cached_tokens_metric.labels().inc.assert_called_once_with( + 100 + ) + + # Verify cache misses metric was NOT called + mock_logger.litellm_cache_misses_metric.labels.assert_not_called() + + def test_increment_cache_metrics_on_cache_miss(self, sample_enum_values): + """Test that cache miss increments the correct metrics""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=False + standard_logging_payload = { + "cache_hit": False, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify cache misses metric was incremented + mock_logger.litellm_cache_misses_metric.labels.assert_called() + mock_logger.litellm_cache_misses_metric.labels().inc.assert_called_once() + + # Verify cache hits and cached tokens metrics were NOT called + mock_logger.litellm_cache_hits_metric.labels.assert_not_called() + mock_logger.litellm_cached_tokens_metric.labels.assert_not_called() + + def test_increment_cache_metrics_when_cache_hit_is_none(self, sample_enum_values): + """Test that no metrics are incremented when cache_hit is None""" + # Create mock for PrometheusLogger instance + mock_logger = MagicMock() + + from litellm.integrations.prometheus import PrometheusLogger + + # Create a mock standard logging payload with cache_hit=None + standard_logging_payload = { + "cache_hit": None, + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + "model_group": "openai", + "request_tags": [], + } + + # Create mock metrics + mock_logger.litellm_cache_hits_metric = MagicMock() + mock_logger.litellm_cache_misses_metric = MagicMock() + mock_logger.litellm_cached_tokens_metric = MagicMock() + mock_logger.get_labels_for_metric = MagicMock( + return_value=[ + "model", + "hashed_api_key", + "api_key_alias", + "team", + "team_alias", + "end_user", + "user", + ] + ) + + # Call the method + PrometheusLogger._increment_cache_metrics( + mock_logger, + standard_logging_payload=standard_logging_payload, + enum_values=sample_enum_values, + ) + + # Verify NO metrics were called + mock_logger.litellm_cache_hits_metric.labels.assert_not_called() + mock_logger.litellm_cache_misses_metric.labels.assert_not_called() + mock_logger.litellm_cached_tokens_metric.labels.assert_not_called() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])