From 13d1caf94a4b654679d6dc4f492b7948075bee40 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 7 Aug 2024 09:43:35 -0700
Subject: [PATCH] add set_remaining_tokens_requests_metric

---
 litellm/integrations/prometheus.py | 33 +++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py
index 4a271d6e004..41679001552 100644
--- a/litellm/integrations/prometheus.py
+++ b/litellm/integrations/prometheus.py
@@ -8,7 +8,7 @@ import subprocess
 import sys
 import traceback
 import uuid
-from typing import Optional, Union
+from typing import Optional, TypedDict, Union
 
 import dotenv
 import requests  # type: ignore
@@ -124,6 +124,29 @@ class PrometheusLogger:
                         "litellm_model_name",
                     ],
                 )
+                # Get all keys
+                _logged_llm_labels = [
+                    "litellm_model_name",
+                    "model_id",
+                    "api_base",
+                    "api_provider",
+                ]
+
+                self.deployment_unhealthy = Gauge(
+                    "deployment_unhealthy",
+                    'Value is "1" when deployment is in an unhealthy state',
+                    labelnames=_logged_llm_labels,
+                )
+                self.deployment_partial_outage = Gauge(
+                    "deployment_partial_outage",
+                    'Value is "1" when deployment is experiencing a partial outage',
+                    labelnames=_logged_llm_labels,
+                )
+                self.deployment_healthy = Gauge(
+                    "deployment_healthy",
+                    'Value is "1" when deployment is in an healthy state',
+                    labelnames=_logged_llm_labels,
+                )
 
         except Exception as e:
             print_verbose(f"Got exception on init prometheus client {str(e)}")
@@ -273,6 +296,7 @@ class PrometheusLogger:
             model_group = _metadata.get("model_group", None)
             api_base = _metadata.get("api_base", None)
             llm_provider = _litellm_params.get("custom_llm_provider", None)
+            model_id = _metadata.get("model_id")
 
             remaining_requests = None
             remaining_tokens = None
@@ -307,6 +331,13 @@ class PrometheusLogger:
                     model_group, llm_provider, api_base, litellm_model_name
                 ).set(remaining_tokens)
 
+            """
+            log these labels
+            ["litellm_model_name", "model_id", "api_base", "api_provider"]
+            """
+            self.deployment_healthy.labels(
+                litellm_model_name, model_id, api_base, llm_provider
+            ).set(1)
         except Exception as e:
             verbose_logger.error(
                 "Prometheus Error: set_remaining_tokens_requests_metric. Exception occured - {}".format(