From 4b7e102187bcc56163bfbfcac38fa85da0dc0e79 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 19 Mar 2024 14:48:38 -0700
Subject: [PATCH] (v0) prometheus metric

---
 litellm/proxy/proxy_server.py | 36 +++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index f718e8901..b8cd1a9fa 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -983,6 +983,42 @@ def cost_tracking():
                 litellm.success_callback.append(_PROXY_track_cost_callback)  # type: ignore
 
 
+from prometheus_client import Counter, REGISTRY
+from prometheus_client import make_asgi_app
+
+# Add prometheus asgi middleware to route /metrics requests
+metrics_app = make_asgi_app()
+app.mount("/metrics", metrics_app)
+try:
+    calls_metric = Counter("calls_metric", "Measure of calls")
+    tokens_metric = Counter("tokens_metric", "Measure of tokens")
+    spend_metric = Counter("spend_metric", "Measure of spend")
+except:
+    pass
+
+# # Define Prometheus counters for metrics
+# Register metrics with the default registry
+
+
+def track_prometheus_metrics(
+    kwargs,  # kwargs to completion
+    completion_response,  # response from completion
+    start_time,
+    end_time,  # start/end time
+):
+    global calls_metric
+    user: str = ("issues",)
+    key: str = ("sk-02Wr4IAlN3NvPXvL5JVvDA",)
+    model: str = ("gpt-3.5-turbo",)
+    budgets = 1
+    tokens = 20
+    # print("incrementing prometheus metrics")
+    calls_metric.labels(user, key, model).inc()
+
+
+litellm.callbacks.append(track_prometheus_metrics)
+
+
 async def _PROXY_track_cost_callback(
     kwargs,  # kwargs to completion
     completion_response: litellm.ModelResponse,  # response from completion