From e09ef4afc72a0a0e2bc34d24984bdd3835d10cc9 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 2 Apr 2025 17:39:48 -0700
Subject: [PATCH] use service logger for tracking pod lock status

---
 litellm/types/services.py | 68 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 1 deletion(-)

diff --git a/litellm/types/services.py b/litellm/types/services.py
index 3eb283dbe9..05944772d7 100644
--- a/litellm/types/services.py
+++ b/litellm/types/services.py
@@ -1,8 +1,15 @@
 import enum
 import uuid
-from typing import Optional
+from typing import List, Optional
 
 from pydantic import BaseModel, Field
+from typing_extensions import TypedDict
+
+
+class ServiceMetrics(enum.Enum):
+    COUNTER = "counter"
+    HISTOGRAM = "histogram"
+    GAUGE = "gauge"
 
 
 class ServiceTypes(str, enum.Enum):
@@ -18,6 +25,62 @@ class ServiceTypes(str, enum.Enum):
     ROUTER = "router"
     AUTH = "auth"
     PROXY_PRE_CALL = "proxy_pre_call"
+    POD_LOCK_MANAGER = "pod_lock_manager"
+
+
+class ServiceConfig(TypedDict):
+    """
+    Configuration for services and their metrics
+    """
+
+    metrics: List[ServiceMetrics]  # What metrics this service should support
+
+
+"""
+Metric types to use for each service
+
+- REDIS only needs Counter, Histogram
+- Pod Lock Manager only needs a gauge metric
+"""
+DEFAULT_SERVICE_CONFIGS = {
+    ServiceTypes.REDIS.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.DB.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.BATCH_WRITE_TO_DB.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.RESET_BUDGET_JOB.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.LITELLM.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.ROUTER.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.AUTH.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.PROXY_PRE_CALL.value: {
+        "metrics": [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
+    },
+    ServiceTypes.POD_LOCK_MANAGER.value: {"metrics": [ServiceMetrics.GAUGE]},
+}
+
+
+class ServiceEventMetadata(TypedDict, total=False):
+    """
+    The metadata logged during service success/failure
+
+    Add any extra fields you expect to access in the service_success_hook/service_failure_hook
+    """
+
+    # Dynamically control gauge labels and values
+    gauge_labels: Optional[str]
+    gauge_value: Optional[float]
 
 
 class ServiceLoggerPayload(BaseModel):
@@ -30,6 +93,9 @@ class ServiceLoggerPayload(BaseModel):
     service: ServiceTypes = Field(description="who is this for? - postgres/redis")
     duration: float = Field(description="How long did the request take?")
     call_type: str = Field(description="The call of the service, being made")
+    event_metadata: Optional[dict] = Field(
+        description="The metadata logged during service success/failure"
+    )
 
     def to_json(self, **kwargs):
         try: