From 1eb6be9d44d27e7fd39ae1088188a08d8e2821e7 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Fri, 17 Apr 2026 17:00:13 +0000 Subject: [PATCH] Removed slots attr which was causing test failures in python 3.9. Improved sentinel handling flagged by greptile --- litellm/integrations/prometheus.py | 23 ++++++----- litellm/types/integrations/prometheus.py | 40 +++---------------- .../test_prometheus_label_value_sanitize.py | 34 ++++++++++++++++ 3 files changed, 52 insertions(+), 45 deletions(-) create mode 100644 tests/test_litellm/types/test_prometheus_label_value_sanitize.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2657aea62c..4995f830d4 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -16,6 +16,7 @@ from typing import ( List, Literal, Optional, + Sequence, Tuple, Union, cast, @@ -38,7 +39,6 @@ from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import ( _sanitize_prometheus_label_name, _sanitize_prometheus_label_value, - _sanitize_prometheus_label_value_v1, ) from litellm.types.utils import StandardLoggingPayload @@ -2659,7 +2659,7 @@ class PrometheusLogger(CustomLogger): self, data_fetch_function: Callable[..., Awaitable[Tuple[List[Any], Optional[int]]]], set_metrics_function: Callable[[List[Any]], Awaitable[None]], - data_type: Literal["teams", "keys", "users"], + data_type: Literal["teams", "keys", "users", "orgs"], ): """ Generic method to initialize budget metrics for teams or API keys. @@ -3493,34 +3493,37 @@ class PrometheusLabelFactoryContext: "_resolved_end_user", ) + _END_USER_NOT_COMPUTED = object() + def __init__(self, enum_values: UserAPIKeyLabelValues) -> None: self.enum_values = enum_values enum_dict = enum_values.model_dump() self._sanitized_enum: Dict[str, Optional[str]] = { - k: _sanitize_prometheus_label_value_v1(v) + k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items() } self._custom_by_sanitized_key: Dict[str, Optional[str]] = {} if enum_values.custom_metadata_labels is not None: for key, value in enum_values.custom_metadata_labels.items(): sk = _sanitize_prometheus_label_name(key) - self._custom_by_sanitized_key[sk] = _sanitize_prometheus_label_value_v1( + self._custom_by_sanitized_key[sk] = _sanitize_prometheus_label_value( value ) self._tag_labels: Dict[str, Optional[str]] = {} if enum_values.tags is not None: for k, v in get_custom_labels_from_tags(enum_values.tags).items(): - self._tag_labels[k] = _sanitize_prometheus_label_value_v1(v) - self._resolved_end_user: Optional[str] = None + self._tag_labels[k] = _sanitize_prometheus_label_value(v) + # Use a dedicated sentinel so `None` can be cached as a computed result. + self._resolved_end_user: Any = self._END_USER_NOT_COMPUTED def get_resolved_end_user(self) -> Optional[str]: - if self._resolved_end_user is None: + if self._resolved_end_user is self._END_USER_NOT_COMPUTED: fn = _get_cached_end_user_id_for_cost_tracking() self._resolved_end_user = fn( litellm_params={"user_api_key_end_user_id": self.enum_values.end_user}, service_type="prometheus", ) - return self._resolved_end_user + return cast(Optional[str], self._resolved_end_user) def _prometheus_labels_from_context( @@ -3646,7 +3649,7 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]: def _tag_matches_wildcard_configured_pattern( - tags: List[str], configured_tag: str + tags: Sequence[str], configured_tag: str ) -> bool: """ Check if any of the request tags matches a wildcard configured pattern @@ -3678,7 +3681,7 @@ def _tag_matches_wildcard_configured_pattern( return any(re.match(pattern=regex_pattern, string=tag) for tag in tags) -def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]: +def get_custom_labels_from_tags(tags: Sequence[str]) -> Dict[str, str]: """ Get custom labels from tags based on admin configuration. diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 6acb5951d4..f13ec2534d 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -2,7 +2,7 @@ import re from dataclasses import dataclass, field, fields from enum import Enum from types import MappingProxyType -from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Tuple +from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Tuple, Union from typing_extensions import Annotated @@ -41,42 +41,11 @@ def _sanitize_prometheus_label_name(label: str) -> str: return sanitized -def _sanitize_prometheus_label_value(value: Optional[Any]) -> Optional[str]: - """ - Sanitize a label value for Prometheus text format compatibility. - - Removes or replaces characters that break the Prometheus exposition format: - - U+2028 (Line Separator) and U+2029 (Paragraph Separator) are removed - - Carriage returns are removed - - Newlines are replaced with spaces - - Backslashes and double quotes are escaped per Prometheus spec - """ - if value is None: - return None - - # Coerce non-string values (int, bool, etc.) to str before sanitizing - str_value: str = value if isinstance(value, str) else str(value) - - # Remove Unicode line/paragraph separators that break text format - str_value = str_value.replace("\u2028", "").replace("\u2029", "") - - # Remove carriage returns - str_value = str_value.replace("\r", "") - - # Replace newlines with spaces - str_value = str_value.replace("\n", " ") - - # Escape backslashes and double quotes per Prometheus exposition format - str_value = str_value.replace("\\", "\\\\").replace('"', '\\"') - - return str_value - - # v1: single translate pass + escape loop (avoids chained str.replace allocations). _PROMETHEUS_LABEL_VALUE_TRANSLATE_V1 = str.maketrans("\n", " ", "\r\u2028\u2029") -def _sanitize_prometheus_label_value_v1(value: Optional[Any]) -> Optional[str]: +def _sanitize_prometheus_label_value(value: Optional[Any]) -> Optional[str]: """ Same semantics as :func:`_sanitize_prometheus_label_value`, implemented with ``str.translate`` plus a single escape pass instead of chained ``replace``. @@ -774,7 +743,7 @@ class PrometheusMetricLabels: return default_labels + custom_labels -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class UserAPIKeyLabelValues: """ Prometheus metric label inputs (Python field names match historical Pydantic ``model_dump`` keys). @@ -794,7 +763,8 @@ class UserAPIKeyLabelValues: requested_model: Optional[str] = None model: Optional[str] = None litellm_model_name: Optional[str] = None - tags: Tuple[str, ...] = () + # Accept list/tuple at construction time; normalize to tuple in __post_init__. + tags: Union[Tuple[str, ...], List[str]] = () custom_metadata_labels: Mapping[str, str] = field(default_factory=dict) model_id: Optional[str] = None api_base: Optional[str] = None diff --git a/tests/test_litellm/types/test_prometheus_label_value_sanitize.py b/tests/test_litellm/types/test_prometheus_label_value_sanitize.py new file mode 100644 index 0000000000..9ff7eb460e --- /dev/null +++ b/tests/test_litellm/types/test_prometheus_label_value_sanitize.py @@ -0,0 +1,34 @@ +import pytest + +from litellm.types.integrations.prometheus import ( + _sanitize_prometheus_label_value, +) + + +@pytest.mark.parametrize( + "value,expected", + [ + (None, None), + ("", ""), + ("plain", "plain"), + # Newlines -> spaces, carriage returns removed + ("a\nb", "a b"), + ("a\rb", "ab"), + ("a\r\nb", "a b"), + # Unicode line/paragraph separators removed + ("a\u2028b", "ab"), + ("a\u2029b", "ab"), + ("a\u2028b\u2029c", "abc"), + # Escapes per Prometheus text format + ('he said "hi"', 'he said \\"hi\\"'), + (r"path\to\file", r"path\\to\\file"), + (r'quote\"slash\\', r'quote\\\"slash\\\\'), + # Non-string inputs get coerced to str first + (123, "123"), + (True, "True"), + (False, "False"), + ], +) +def test_sanitize_prometheus_label_value_expected_outputs(value, expected): + assert _sanitize_prometheus_label_value(value) == expected +