phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(prometheus.py): add user tracking to prometheus

Browse source
This commit is contained in:
Krrish Dholakia 2024-04-22 15:14:38 -07:00
parent 14d1bd97d4
commit 49852b68c1
1 changed files with 11 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

19
litellm/integrations/prometheus.py
View file

@ -25,27 +25,27 @@ class PrometheusLogger:
self.litellm_llm_api_failed_requests_metric = Counter( self.litellm_llm_api_failed_requests_metric = Counter(
name="litellm_llm_api_failed_requests_metric", name="litellm_llm_api_failed_requests_metric",
documentation="Total number of failed LLM API calls via litellm", documentation="Total number of failed LLM API calls via litellm",
labelnames=["end_user", "hashed_api_key", "model", "team"], labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
) )
self.litellm_requests_metric = Counter( self.litellm_requests_metric = Counter(
name="litellm_requests_metric", name="litellm_requests_metric",
documentation="Total number of LLM calls to litellm", documentation="Total number of LLM calls to litellm",
labelnames=["end_user", "hashed_api_key", "model", "team"], labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
) )
# Counter for spend # Counter for spend
self.litellm_spend_metric = Counter( self.litellm_spend_metric = Counter(
"litellm_spend_metric", "litellm_spend_metric",
"Total spend on LLM requests", "Total spend on LLM requests",
labelnames=["end_user", "hashed_api_key", "model", "team"], labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
) )
# Counter for total_output_tokens # Counter for total_output_tokens
self.litellm_tokens_metric = Counter( self.litellm_tokens_metric = Counter(
"litellm_total_tokens", "litellm_total_tokens",
"Total number of input + output tokens from LLM requests", "Total number of input + output tokens from LLM requests",
labelnames=["end_user", "hashed_api_key", "model", "team"], labelnames=["end_user", "hashed_api_key", "model", "team", "user"],
) )
except Exception as e: except Exception as e:
print_verbose(f"Got exception on init prometheus client {str(e)}") print_verbose(f"Got exception on init prometheus client {str(e)}")
@ -71,6 +71,9 @@ class PrometheusLogger:
litellm_params = kwargs.get("litellm_params", {}) or {} litellm_params = kwargs.get("litellm_params", {}) or {}
proxy_server_request = litellm_params.get("proxy_server_request") or {} proxy_server_request = litellm_params.get("proxy_server_request") or {}
end_user_id = proxy_server_request.get("body", {}).get("user", None) end_user_id = proxy_server_request.get("body", {}).get("user", None)
user_id = proxy_server_request.get("metadata", {}).get(
"user_api_key_user_id", None
)
user_api_key = litellm_params.get("metadata", {}).get("user_api_key", None) user_api_key = litellm_params.get("metadata", {}).get("user_api_key", None)
user_api_team = litellm_params.get("metadata", {}).get( user_api_team = litellm_params.get("metadata", {}).get(
"user_api_key_team_id", None "user_api_key_team_id", None
@ -94,19 +97,19 @@ class PrometheusLogger:
user_api_key = hash_token(user_api_key) user_api_key = hash_token(user_api_key)
self.litellm_requests_metric.labels( self.litellm_requests_metric.labels(
end_user_id, user_api_key, model, user_api_team end_user_id, user_api_key, model, user_api_team, user_id
).inc() ).inc()
self.litellm_spend_metric.labels( self.litellm_spend_metric.labels(
end_user_id, user_api_key, model, user_api_team end_user_id, user_api_key, model, user_api_team, user_id
).inc(response_cost) ).inc(response_cost)
self.litellm_tokens_metric.labels( self.litellm_tokens_metric.labels(
end_user_id, user_api_key, model, user_api_team end_user_id, user_api_key, model, user_api_team, user_id
).inc(tokens_used) ).inc(tokens_used)
### FAILURE INCREMENT ### ### FAILURE INCREMENT ###
if "exception" in kwargs: if "exception" in kwargs:
self.litellm_llm_api_failed_requests_metric.labels( self.litellm_llm_api_failed_requests_metric.labels(
end_user_id, user_api_key, model, user_api_team end_user_id, user_api_key, model, user_api_team, user_id
).inc() ).inc()
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()