From 270dad93d5800660d341f17eaa31beaadff1d0c0 Mon Sep 17 00:00:00 2001
From: Bram van Meurs <7422223+bcvanmeurs@users.noreply.github.com>
Date: Thu, 13 Jun 2024 12:50:40 +0200
Subject: [PATCH] feat(prometheus): add api_team_alias to exported labels

---
 litellm/integrations/prometheus.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py
index af0d1d310b0..c97fce0d837 100644
--- a/litellm/integrations/prometheus.py
+++ b/litellm/integrations/prometheus.py
@@ -22,27 +22,27 @@ class PrometheusLogger:
             self.litellm_llm_api_failed_requests_metric = Counter(
                 name="litellm_llm_api_failed_requests_metric",
                 documentation="Total number of failed LLM API calls via litellm",
-                labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
+                labelnames=["end_user", "hashed_api_key", "model", "team", "team_alias", "user"],
             )
 
             self.litellm_requests_metric = Counter(
                 name="litellm_requests_metric",
                 documentation="Total number of LLM calls to litellm",
-                labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
+                labelnames=["end_user", "hashed_api_key", "model", "team", "team_alias", "user"],
             )
 
             # Counter for spend
             self.litellm_spend_metric = Counter(
                 "litellm_spend_metric",
                 "Total spend on LLM requests",
-                labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
+                labelnames=["end_user", "hashed_api_key", "model", "team", "team_alias", "user"],
             )
 
             # Counter for total_output_tokens
             self.litellm_tokens_metric = Counter(
                 "litellm_total_tokens",
                 "Total number of input + output tokens from LLM requests",
-                labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
+                labelnames=["end_user", "hashed_api_key", "model", "team", "team_alias", "user"],
             )
         except Exception as e:
             print_verbose(f"Got exception on init prometheus client {str(e)}")
@@ -75,6 +75,9 @@ class PrometheusLogger:
             user_api_team = litellm_params.get("metadata", {}).get(
                 "user_api_key_team_id", None
             )
+            user_api_team_alias = litellm_params.get("metadata", {}).get(
+                "user_api_key_team_alias", None
+            )
             if response_obj is not None:
                 tokens_used = response_obj.get("usage", {}).get("total_tokens", 0)
             else:
@@ -94,19 +97,19 @@ class PrometheusLogger:
                 user_api_key = hash_token(user_api_key)
 
             self.litellm_requests_metric.labels(
-                end_user_id, user_api_key, model, user_api_team, user_id
+                end_user_id, user_api_key, model, user_api_team, user_api_team_alias, user_id
             ).inc()
             self.litellm_spend_metric.labels(
-                end_user_id, user_api_key, model, user_api_team, user_id
+                end_user_id, user_api_key, model, user_api_team, user_api_team_alias, user_id
             ).inc(response_cost)
             self.litellm_tokens_metric.labels(
-                end_user_id, user_api_key, model, user_api_team, user_id
+                end_user_id, user_api_key, model, user_api_team, user_api_team_alias, user_id
             ).inc(tokens_used)
 
             ### FAILURE INCREMENT ###
             if "exception" in kwargs:
                 self.litellm_llm_api_failed_requests_metric.labels(
-                    end_user_id, user_api_key, model, user_api_team, user_id
+                    end_user_id, user_api_key, model, user_api_team, user_api_team_alias, user_id
                 ).inc()
         except Exception as e:
             verbose_logger.error(