feat(prometheus): add caching metrics (#18755)

This commit is contained in:
Harshit Jain 2026-01-07 21:29:04 +05:30 • committed by GitHub
parent 91b5c66cf2
commit 92f7789f10
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 318 additions and 15 deletions

View file

@ -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",

View file

@ -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)

View file

@ -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"])