From 06a362d35fb25c4c0ab7bd239fb765a50c48392e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 22 Aug 2024 13:58:10 -0700
Subject: [PATCH] track litellm_request_latency_metric

---
 litellm/integrations/prometheus.py | 38 ++++++++++++++++++++++++++++++
 litellm/proxy/proxy_config.yaml    |  4 +++-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py
index 1471f59b75c..dadafa80e35 100644
--- a/litellm/integrations/prometheus.py
+++ b/litellm/integrations/prometheus.py
@@ -60,6 +60,25 @@ class PrometheusLogger(CustomLogger):
                 ],
             )
 
+            # request latency metrics
+            self.litellm_request_latency_metric = Histogram(
+                "litellm_request_latency_metric",
+                "Total latency (seconds) for a request to LiteLLM",
+                labelnames=[
+                    "model",
+                    "litellm_call_id",
+                ],
+            )
+
+            self.litellm_deployment_latency_metric = Histogram(
+                "litellm_deployment_latency_metric",
+                "Total latency (seconds) for a models LLM API call",
+                labelnames=[
+                    "model",
+                    "litellm_call_id",
+                ],
+            )
+
             # Counter for spend
             self.litellm_spend_metric = Counter(
                 "litellm_spend_metric",
@@ -329,6 +348,25 @@ class PrometheusLogger(CustomLogger):
             user_api_key, user_api_key_alias, model_group
         ).set(remaining_tokens)
 
+        # latency metrics
+        total_time: timedelta = kwargs.get("end_time") - kwargs.get("start_time")
+        total_time_seconds = total_time.total_seconds()
+        api_call_total_time: timedelta = kwargs.get("end_time") - kwargs.get(
+            "api_call_start_time"
+        )
+
+        api_call_total_time_seconds = api_call_total_time.total_seconds()
+
+        litellm_call_id = kwargs.get("litellm_call_id")
+
+        self.litellm_request_latency_metric.labels(model, litellm_call_id).observe(
+            total_time_seconds
+        )
+
+        self.litellm_deployment_latency_metric.labels(model, litellm_call_id).observe(
+            api_call_total_time_seconds
+        )
+
         # set x-ratelimit headers
         if premium_user is True:
             self.set_llm_deployment_success_metrics(
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 65c7f70525c..7c524eb18a0 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -4,7 +4,9 @@ model_list:
       model: openai/fake
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
-
+litellm_settings:
+  success_callback: ["prometheus"]
+  failure_callback: ["prometheus"]
 guardrails:
   - guardrail_name: "lakera-pre-guard"
     litellm_params: