From 412d30d3627d2ad1811f5c9c2e1ee9c4c8d1e647 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 17 Aug 2024 10:02:20 -0700
Subject: [PATCH] add litellm-key-remaining-tokens on prometheus

---
 litellm/integrations/prometheus.py            | 49 +++++++++++++++++--
 .../proxy/hooks/parallel_request_limiter.py   | 11 +++++
 litellm/proxy/proxy_config.yaml               |  4 +-
 3 files changed, 56 insertions(+), 8 deletions(-)

diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py
index 08431fd7af9..8ee726196ba 100644
--- a/litellm/integrations/prometheus.py
+++ b/litellm/integrations/prometheus.py
@@ -103,13 +103,30 @@ class PrometheusLogger(CustomLogger):
                 "Remaining budget for api key",
                 labelnames=["hashed_api_key", "api_key_alias"],
             )
-
-            ########################################
-            # LLM API Deployment Metrics / analytics
-            ########################################
-
             # Litellm-Enterprise Metrics
             if premium_user is True:
+
+                ########################################
+                # LiteLLM Virtual API KEY metrics
+                ########################################
+                # Remaining MODEL RPM limit for API Key
+                self.litellm_remaining_api_key_requests_for_model = Gauge(
+                    "litellm_remaining_api_key_requests_for_model",
+                    "Remaining Requests API Key can make for model (model based rpm limit on key)",
+                    labelnames=["hashed_api_key", "api_key_alias", "model"],
+                )
+
+                # Remaining MODEL TPM limit for API Key
+                self.litellm_remaining_api_key_tokens_for_model = Gauge(
+                    "litellm_remaining_api_key_tokens_for_model",
+                    "Remaining Tokens API Key can make for model (model based tpm limit on key)",
+                    labelnames=["hashed_api_key", "api_key_alias", "model"],
+                )
+
+                ########################################
+                # LLM API Deployment Metrics / analytics
+                ########################################
+
                 # Remaining Rate Limit for model
                 self.litellm_remaining_requests_metric = Gauge(
                     "litellm_remaining_requests",
@@ -197,6 +214,7 @@ class PrometheusLogger(CustomLogger):
         model = kwargs.get("model", "")
         response_cost = kwargs.get("response_cost", 0.0) or 0
         litellm_params = kwargs.get("litellm_params", {}) or {}
+        _metadata = litellm_params.get("metadata", {})
         proxy_server_request = litellm_params.get("proxy_server_request") or {}
         end_user_id = proxy_server_request.get("body", {}).get("user", None)
         user_id = litellm_params.get("metadata", {}).get("user_api_key_user_id", None)
@@ -286,6 +304,27 @@ class PrometheusLogger(CustomLogger):
             user_api_key, user_api_key_alias
         ).set(_remaining_api_key_budget)
 
+        # Set remaining rpm/tpm for API Key + model
+        # see parallel_request_limiter.py - variables are set there
+        model_group = _metadata.get("model_group", None)
+        remaining_requests_variable_name = (
+            f"litellm-key-remaining-requests-{model_group}"
+        )
+        remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
+
+        remaining_requests = _metadata.get(
+            remaining_requests_variable_name, sys.maxsize
+        )
+        remaining_tokens = _metadata.get(remaining_tokens_variable_name, sys.maxsize)
+
+        self.litellm_remaining_api_key_requests_for_model.labels(
+            user_api_key, user_api_key_alias, model
+        ).set(remaining_requests)
+
+        self.litellm_remaining_api_key_tokens_for_model.labels(
+            user_api_key, user_api_key_alias, model
+        ).set(remaining_tokens)
+
         # set x-ratelimit headers
         if premium_user is True:
             self.set_llm_deployment_success_metrics(
diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py
index 4af86c24c43..b7f923631e7 100644
--- a/litellm/proxy/hooks/parallel_request_limiter.py
+++ b/litellm/proxy/hooks/parallel_request_limiter.py
@@ -263,6 +263,17 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
                         request_count_api_key, new_val
                     )
 
+            tpm_limit_for_model = tpm_limit_for_model or sys.maxsize
+            rpm_limit_for_model = rpm_limit_for_model or sys.maxsize
+            # Add remaining tokens, requests to metadata
+            _remaining_tokens = tpm_limit_for_model - new_val["current_tpm"]
+            _remaining_requests = rpm_limit_for_model - new_val["current_rpm"]
+            _remaining_limits_data = {
+                f"litellm-key-remaining-tokens-{_model}": _remaining_tokens,
+                f"litellm-key-remaining-requests-{_model}": _remaining_requests,
+            }
+            data["metadata"].update(_remaining_limits_data)
+
         # check if REQUEST ALLOWED for user_id
         user_id = user_api_key_dict.user_id
         if user_id is not None:
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index d25f1b94682..336cd8458a6 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -42,7 +42,5 @@ general_settings:
 
 litellm_settings:
   fallbacks: [{"gemini-1.5-pro-001": ["gpt-4o"]}]
-  callbacks: ["gcs_bucket"]
-  success_callback: ["langfuse"]
+  success_callback: ["langfuse", "prometheus"]
   langfuse_default_tags: ["cache_hit", "cache_key", "user_api_key_alias", "user_api_key_team_alias"]
-  cache: True