From 8d7934c4ba77449a9bb6b38bb0d6e17c79b60030 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 5 Nov 2024 13:12:14 -0800
Subject: [PATCH] track ALLOW_FAILED_DB_REQUESTS on prometheus

---
 litellm/proxy/auth/auth_checks.py | 25 ++++++++++++++++++++++---
 litellm/proxy/proxy_config.yaml   |  6 +++---
 litellm/types/services.py         |  1 +
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index e8b27966fbb..9e535a3eb2a 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -719,14 +719,16 @@ async def get_key_object(
 
         return _response
     except httpx.ConnectError as e:
-        return _handle_failed_db_connection_for_get_key_object(e=e)
+        return await _handle_failed_db_connection_for_get_key_object(e=e)
     except Exception as e:
         raise Exception(
             f"Key doesn't exist in db. key={hashed_token}. Create key via `/key/generate` call."
         )
 
 
-def _handle_failed_db_connection_for_get_key_object(e: Exception):
+async def _handle_failed_db_connection_for_get_key_object(
+    e: Exception,
+) -> UserAPIKeyAuth:
     """
     Handles httpx.ConnectError when reading a Virtual Key from LiteLLM DB
 
@@ -738,9 +740,26 @@ def _handle_failed_db_connection_for_get_key_object(e: Exception):
     Raises:
         - Orignal Exception in all other cases
     """
-    from litellm.proxy.proxy_server import general_settings
+    from litellm.proxy.proxy_server import general_settings, proxy_logging_obj
 
+    # If this flag is on, requests failing to connect to the DB will be allowed
     if general_settings.get("allow_failed_db_requests", True):
+        # log to prometheus
+        await proxy_logging_obj.service_logging_obj.async_service_failure_hook(
+            error=e,
+            service=ServiceTypes.ALLOW_FAILED_DB_REQUESTS,
+            call_type="get_key_object",
+            parent_otel_span=None,
+            duration=0.0,
+            start_time=None,
+            end_time=None,
+            event_metadata={
+                "function_name": "get_key_object",
+                "function_kwargs": {},
+                "function_args": [],
+            },
+        )
+
         return UserAPIKeyAuth(
             key_name="failed-to-connect-to-db", token="failed-to-connect-to-db"
         )
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 23834f75970..d79bcdb0341 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -6,6 +6,6 @@ model_list:
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
 
 
-general_settings: 
-    alerting: ["slack"]
-    alerting_threshold: 0.001
+litellm_settings:
+  callbacks: ["prometheus"]
+  service_callback: ["prometheus_system"]
diff --git a/litellm/types/services.py b/litellm/types/services.py
index cfa427ebc3d..161832404c1 100644
--- a/litellm/types/services.py
+++ b/litellm/types/services.py
@@ -12,6 +12,7 @@ class ServiceTypes(str, enum.Enum):
 
     REDIS = "redis"
     DB = "postgres"
+    ALLOW_FAILED_DB_REQUESTS = "allow_failed_db_requests"
     BATCH_WRITE_TO_DB = "batch_write_to_db"
     LITELLM = "self"
     ROUTER = "router"