From aaf6afff9b4d94aa6ffd8fc087e34d3ee3f57082 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Mon, 29 Jun 2026 12:35:31 +0300 Subject: [PATCH] feat(prometheus): add litellm_overhead_with_guardrails_latency_metric litellm_overhead_latency_metric only covers the SDK wrapper window and excludes proxy guardrails. Add a histogram that sums SDK overhead plus pre/post-call guardrail durations (during-call excluded since it runs concurrently with the LLM call, alongside logging_only and MCP modes that never block the response), recorded next to the existing overhead metric with the same labels and buckets. The name spells out "with guardrails" so it reads distinctly from the existing litellm_overhead_latency_metric rather than as a vague second overhead number. guardrail_information is typed as a list but some guardrails assign a single dict directly; normalize that shape to a one-item list so the success-metrics block doesn't crash with AttributeError. No existing metric's value is changed. --- litellm/integrations/prometheus.py | 83 +++++- litellm/types/integrations/prometheus.py | 11 + ...est_prometheus_overhead_with_guardrails.py | 238 ++++++++++++++++++ 3 files changed, 331 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fcec551f25aa..1f516e9dc935 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -49,12 +49,16 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import ( _sanitize_prometheus_label_name, _sanitize_prometheus_label_value, ) -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ( + StandardLoggingGuardrailInformation, + StandardLoggingPayload, +) if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler @@ -65,6 +69,8 @@ class PrometheusLogger(CustomLogger): # Class variables or attributes + _ADDITIVE_GUARDRAIL_MODES = frozenset((GuardrailEventHooks.pre_call.value, GuardrailEventHooks.post_call.value)) + @staticmethod def get_instance() -> Optional["PrometheusLogger"]: """Find the PrometheusLogger instance from litellm.callbacks, if registered.""" @@ -343,6 +349,14 @@ def __init__( buckets=self.latency_buckets, ) + self.litellm_overhead_with_guardrails_latency_metric = self._histogram_factory( + "litellm_overhead_with_guardrails_latency_metric", + "Total internal latency (seconds) added by LiteLLM, including " + "pre/post-call guardrails (excludes the LLM API call)", + labelnames=self.get_labels_for_metric("litellm_overhead_with_guardrails_latency_metric"), + buckets=self.latency_buckets, + ) + # Request queue time metric self.litellm_request_queue_time_metric = self._histogram_factory( "litellm_request_queue_time_seconds", @@ -1001,6 +1015,67 @@ def get_labels_for_metric(self, metric_name: DEFINED_PROMETHEUS_METRICS) -> List self._cached_metric_labels[metric_name] = filtered_labels return filtered_labels + @staticmethod + def _guardrail_is_additive(info: StandardLoggingGuardrailInformation) -> bool: + mode = info.get("guardrail_mode") + modes = mode if isinstance(mode, list) else [mode] + mode_values = frozenset( + m.value if isinstance(m, GuardrailEventHooks) else m for m in modes if isinstance(m, str) + ) + return bool(mode_values) and mode_values <= PrometheusLogger._ADDITIVE_GUARDRAIL_MODES + + @staticmethod + def _get_guardrail_overhead_seconds( + standard_logging_payload: StandardLoggingPayload, + ) -> float: + """Seconds of additive guardrail time (pre/post-call only) on the payload. + + during_call guardrails run concurrently with the LLM call, so their + wall-clock overlaps the provider call and is not additive overhead; + logging_only and MCP modes never block the user-facing response. A + guardrail counts only when every mode it carries is pre/post-call, so a + mixed list such as ["pre_call", "during_call"] is excluded. + + guardrail_information is typed as a list, but some guardrails assign a + single dict directly, so normalize that shape to a one-item list. + """ + guardrail_information = standard_logging_payload.get("guardrail_information") + entries: list[StandardLoggingGuardrailInformation] = ( + [cast("StandardLoggingGuardrailInformation", guardrail_information)] + if isinstance(guardrail_information, dict) + else guardrail_information or [] + ) + return sum( + (float(info.get("duration") or 0.0) for info in entries if PrometheusLogger._guardrail_is_additive(info)), + 0.0, + ) + + def _set_overhead_with_guardrails_metric( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + label_context: Optional[PrometheusLabelFactoryContext] = None, + ) -> None: + """Record litellm_overhead_with_guardrails_latency_metric (seconds): SDK overhead + + pre/post-call guardrail time. Recorded outside the SDK-overhead gate so + guardrail-only overhead is still captured when litellm_overhead_time_ms + is 0 or absent. + """ + litellm_overhead_time_ms = standard_logging_payload["hidden_params"].get("litellm_overhead_time_ms") + guardrail_overhead_seconds = self._get_guardrail_overhead_seconds(standard_logging_payload) + if litellm_overhead_time_ms is None and guardrail_overhead_seconds <= 0: + return + labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_overhead_with_guardrails_latency_metric" + ), + enum_values=enum_values, + label_context=label_context, + ) + self.litellm_overhead_with_guardrails_latency_metric.labels(**labels).observe( + ((litellm_overhead_time_ms or 0.0) / 1000) + guardrail_overhead_seconds + ) + def _track_end_user_metric_series( self, metric: Any, @@ -2346,6 +2421,12 @@ def set_llm_deployment_success_metrics( litellm_overhead_time_ms / 1000 ) # set as seconds + self._set_overhead_with_guardrails_metric( + standard_logging_payload=standard_logging_payload, + enum_values=enum_values, + label_context=label_context, + ) + if remaining_requests: """ "model_group", diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8f460b799550..fca3319254cc 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -195,6 +195,7 @@ class UserAPIKeyLabelNames(Enum): "litellm_llm_api_time_to_first_token_metric", "litellm_request_total_latency_metric", "litellm_overhead_latency_metric", + "litellm_overhead_with_guardrails_latency_metric", "litellm_remaining_requests_metric", "litellm_remaining_tokens_metric", "litellm_proxy_total_requests_metric", @@ -379,6 +380,16 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.MODEL_ID.value, ] + litellm_overhead_with_guardrails_latency_metric = [ + UserAPIKeyLabelNames.MODEL_GROUP.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.MODEL_ID.value, + ] + litellm_remaining_requests_metric = [ UserAPIKeyLabelNames.MODEL_GROUP.value, UserAPIKeyLabelNames.API_PROVIDER.value, diff --git a/tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py b/tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py new file mode 100644 index 000000000000..56c9702189e4 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_overhead_with_guardrails.py @@ -0,0 +1,238 @@ +""" +Unit tests for litellm_overhead_with_guardrails_latency_metric. + +The metric reports total internal latency LiteLLM adds around the provider +call = SDK overhead (litellm_overhead_time_ms) + pre/post-call guardrail +durations. During-call (moderation) guardrails run concurrently with the LLM +call and are excluded so they don't inflate the overhead. +""" + +from unittest.mock import MagicMock + +import pytest +from prometheus_client import REGISTRY + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import StandardLoggingPayload + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + """Clean up prometheus registry before/after each test.""" + for collector in list(REGISTRY._collector_to_names.keys()): + REGISTRY.unregister(collector) + yield + for collector in list(REGISTRY._collector_to_names.keys()): + REGISTRY.unregister(collector) + + +def test_get_guardrail_overhead_seconds_sums_pre_post_excludes_during(): + """Helper sums pre/post durations, excludes during_call, tolerates missing values.""" + payload = StandardLoggingPayload( + guardrail_information=[ + {"guardrail_mode": GuardrailEventHooks.pre_call, "duration": 0.1}, + {"guardrail_mode": GuardrailEventHooks.during_call, "duration": 0.5}, + {"guardrail_mode": GuardrailEventHooks.post_call, "duration": 0.25}, + {"guardrail_mode": GuardrailEventHooks.post_call}, # no duration -> 0 + ], + ) + # 0.1 (pre) + 0.25 (post) = 0.35; during_call 0.5 excluded; missing duration -> 0 + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.35) < 1e-6 + + +def test_get_guardrail_overhead_seconds_no_guardrails_is_zero(): + """No guardrail_information at all -> 0.0.""" + assert ( + PrometheusLogger._get_guardrail_overhead_seconds( + StandardLoggingPayload(model="gpt-4o") + ) + == 0.0 + ) + + +def test_get_guardrail_overhead_seconds_accepts_plain_string_mode(): + """guardrail_mode may arrive as a plain string after serialization.""" + payload = StandardLoggingPayload( + guardrail_information=[ + {"guardrail_mode": "pre_call", "duration": 0.2}, + {"guardrail_mode": "during_call", "duration": 0.9}, + ], + ) + # only pre_call counts; during_call excluded + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.2) < 1e-6 + + +def test_get_guardrail_overhead_seconds_excludes_list_mode_with_during_call(): + """A list-typed guardrail_mode containing during_call must be excluded. + + guardrail_mode is typed Optional[Union[GuardrailEventHooks, + List[GuardrailEventHooks], GuardrailMode]]; a list mixing in during_call is + not additive (concurrent) overhead and must not be counted. + """ + payload = StandardLoggingPayload( + guardrail_information=[ + { + "guardrail_mode": [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + ], + "duration": 0.3, + }, + {"guardrail_mode": GuardrailEventHooks.post_call, "duration": 0.05}, + ], + ) + # the list entry mixes in during_call -> excluded; only the post_call counts + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.05) < 1e-6 + + +def test_get_guardrail_overhead_seconds_counts_pure_pre_post_list_mode(): + """A list-typed mode containing only additive (pre/post) phases is counted.""" + payload = StandardLoggingPayload( + guardrail_information=[ + {"guardrail_mode": [GuardrailEventHooks.pre_call], "duration": 0.1}, + {"guardrail_mode": ["post_call"], "duration": 0.2}, + ], + ) + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.3) < 1e-6 + + +def test_get_guardrail_overhead_seconds_excludes_logging_only_and_mcp(): + """logging_only and MCP-specific modes do not block the response -> excluded.""" + payload = StandardLoggingPayload( + guardrail_information=[ + {"guardrail_mode": GuardrailEventHooks.logging_only, "duration": 0.4}, + {"guardrail_mode": GuardrailEventHooks.pre_mcp_call, "duration": 0.3}, + {"guardrail_mode": GuardrailEventHooks.during_mcp_call, "duration": 0.2}, + {"guardrail_mode": GuardrailEventHooks.post_call, "duration": 0.05}, + ], + ) + # only the post_call guardrail is additive, user-visible overhead + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.05) < 1e-6 + + +def test_get_guardrail_overhead_seconds_ignores_dict_mode_without_error(): + """guardrail_mode may be a GuardrailMode TypedDict (an unhashable dict at + runtime, from the enterprise Mode-hook path). It must not raise TypeError and + must not be counted (the phase can't be resolved to a blocking pre/post).""" + payload = StandardLoggingPayload( + guardrail_information=[ + # GuardrailMode TypedDict -> plain dict at runtime + {"guardrail_mode": {"tags": {"default": ["pre_call"]}}, "duration": 0.3}, + {"guardrail_mode": GuardrailEventHooks.post_call, "duration": 0.05}, + ], + ) + # dict-typed mode is ignored (no TypeError); only the post_call entry counts + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.05) < 1e-6 + + +def test_get_guardrail_overhead_seconds_ignores_dict_inside_list_mode(): + """A list-typed mode containing a dict must not raise and the dict is ignored.""" + payload = StandardLoggingPayload( + guardrail_information=[ + {"guardrail_mode": [GuardrailEventHooks.pre_call, {"k": "v"}], "duration": 0.1}, + ], + ) + # the dict is ignored; remaining mode is pre_call -> counted + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.1) < 1e-6 + + +def test_get_guardrail_overhead_seconds_handles_single_dict_payload(): + """guardrail_information may be a single dict (e.g. xecguard) rather than a + list. Iterating it would yield string keys and crash the success-metrics + block, so a lone dict must be evaluated as one entry, not raise.""" + payload = StandardLoggingPayload( + guardrail_information={ + "guardrail_mode": "logging_only", + "duration": 0.7, + "guardrail_name": "xecguard", + }, + ) + # the single dict is logging_only -> excluded, and must not raise + assert PrometheusLogger._get_guardrail_overhead_seconds(payload) == 0.0 + + +def test_get_guardrail_overhead_seconds_counts_single_pre_call_dict(): + """A single pre/post-call dict (not wrapped in a list) is still counted.""" + payload = StandardLoggingPayload( + guardrail_information={"guardrail_mode": "pre_call", "duration": 0.3}, + ) + assert abs(PrometheusLogger._get_guardrail_overhead_seconds(payload) - 0.3) < 1e-6 + + +def _patch_label_factory(monkeypatch): + monkeypatch.setattr( + "litellm.integrations.prometheus.prometheus_label_factory", + lambda **kwargs: {}, + ) + + +def test_overhead_with_guardrails_recorded_when_only_guardrails_no_sdk_overhead(monkeypatch): + """Guardrail-only overhead is recorded even when SDK overhead is absent.""" + _patch_label_factory(monkeypatch) + logger = PrometheusLogger() + mock_metric = MagicMock() + logger.litellm_overhead_with_guardrails_latency_metric = mock_metric + + payload = StandardLoggingPayload( + hidden_params={}, # no litellm_overhead_time_ms + guardrail_information=[ + {"guardrail_mode": GuardrailEventHooks.post_call, "duration": 0.2} + ], + ) + logger._set_overhead_with_guardrails_metric( + payload, enum_values=MagicMock(), label_context=MagicMock() + ) + + mock_metric.labels.return_value.observe.assert_called_once() + observed = mock_metric.labels.return_value.observe.call_args[0][0] + assert abs(observed - 0.2) < 1e-6 + + +def test_overhead_with_guardrails_recorded_when_sdk_overhead_is_zero(monkeypatch): + """SDK overhead of exactly 0 (walrus-falsy) must not suppress the metric.""" + _patch_label_factory(monkeypatch) + logger = PrometheusLogger() + mock_metric = MagicMock() + logger.litellm_overhead_with_guardrails_latency_metric = mock_metric + + payload = StandardLoggingPayload( + hidden_params={"litellm_overhead_time_ms": 0.0}, + guardrail_information=[ + {"guardrail_mode": GuardrailEventHooks.pre_call, "duration": 0.1} + ], + ) + logger._set_overhead_with_guardrails_metric( + payload, enum_values=MagicMock(), label_context=MagicMock() + ) + + observed = mock_metric.labels.return_value.observe.call_args[0][0] + assert abs(observed - 0.1) < 1e-6 + + +def test_overhead_with_guardrails_skipped_when_no_overhead_and_no_guardrails(monkeypatch): + """Nothing to record -> the metric is not touched.""" + _patch_label_factory(monkeypatch) + logger = PrometheusLogger() + mock_metric = MagicMock() + logger.litellm_overhead_with_guardrails_latency_metric = mock_metric + + payload = StandardLoggingPayload(hidden_params={}) + logger._set_overhead_with_guardrails_metric( + payload, enum_values=MagicMock(), label_context=MagicMock() + ) + + mock_metric.labels.assert_not_called() + + +def test_overhead_with_guardrails_metric_is_registered(): + """The overhead-with-guardrails histogram is defined and registered on logger init.""" + logger = PrometheusLogger() + assert logger.litellm_overhead_with_guardrails_latency_metric is not None + + registered = [ + name + for name in REGISTRY._names_to_collectors + if name.startswith("litellm_overhead_with_guardrails_latency_metric") + ] + assert registered, "litellm_overhead_with_guardrails_latency_metric not registered"