From bcf42fd82d8fb5806cced0287f65b05bf779117b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 2 Apr 2025 21:19:05 -0700 Subject: [PATCH] linting fix prometheus services --- litellm/integrations/prometheus_services.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index dddaa4d0647..37f0d696fbc 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -3,7 +3,7 @@ # On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers) -from typing import List, Optional, Union +from typing import Dict, List, Optional, Union from litellm._logging import print_verbose, verbose_logger from litellm.types.integrations.prometheus import LATENCY_BUCKETS @@ -43,7 +43,9 @@ class PrometheusServicesLogger: verbose_logger.debug("in init prometheus services metrics") self.services = [item for item in ServiceTypes] - self.payload_to_prometheus_map = {} + self.payload_to_prometheus_map: Dict[ + str, List[Union[Histogram, Counter, Gauge, Collector]] + ] = {} for service in self.services: service_metrics: List[Union[Histogram, Counter, Gauge, Collector]] = [] @@ -78,7 +80,7 @@ class PrometheusServicesLogger: service_metrics.append(gauge) if service_metrics: - self.payload_to_prometheus_map[service] = service_metrics + self.payload_to_prometheus_map[service.value] = service_metrics self.prometheus_to_amount_map: dict = {} ### MOCK TESTING ###