From 039617766f613bf89e26b5cd93bd84a308d23f7a Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 15 Feb 2024 20:58:39 -0800 Subject: [PATCH] (feat) track team_id in langfuse tags --- litellm/integrations/langfuse.py | 1 + litellm/proxy/proxy_server.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py index b548f22bc4..7ce843339a 100644 --- a/litellm/integrations/langfuse.py +++ b/litellm/integrations/langfuse.py @@ -259,6 +259,7 @@ class LangFuseLogger: if key in [ "user_api_key", "user_api_key_user_id", + "user_api_key_team_id", "semantic-similarity", ]: tags.append(f"{key}:{value}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2931851910..0746b4e108 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -819,6 +819,7 @@ async def _PROXY_track_cost_callback( user_id = user_id or kwargs["litellm_params"]["metadata"].get( "user_api_key_user_id", None ) + team_id = kwargs["litellm_params"]["metadata"].get("user_api_key_team_id", None) if kwargs.get("response_cost", None) is not None: response_cost = kwargs["response_cost"] user_api_key = kwargs["litellm_params"]["metadata"].get( @@ -2134,6 +2135,9 @@ async def completion( data["metadata"]["user_api_key"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_team_id"] = getattr( + user_api_key_dict, "team_id", None + ) _headers = dict(request.headers) _headers.pop( "authorization", None @@ -2297,6 +2301,9 @@ async def chat_completion( data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_team_id"] = getattr( + user_api_key_dict, "team_id", None + ) data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata _headers = dict(request.headers) _headers.pop( @@ -2518,6 +2525,9 @@ async def embeddings( ) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_team_id"] = getattr( + user_api_key_dict, "team_id", None + ) data["metadata"]["endpoint"] = str(request.url) ### TEAM-SPECIFIC PARAMS ### @@ -2689,6 +2699,9 @@ async def image_generation( ) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_team_id"] = getattr( + user_api_key_dict, "team_id", None + ) data["metadata"]["endpoint"] = str(request.url) ### TEAM-SPECIFIC PARAMS ### @@ -2844,6 +2857,9 @@ async def moderations( ) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_team_id"] = getattr( + user_api_key_dict, "team_id", None + ) data["metadata"]["endpoint"] = str(request.url) ### TEAM-SPECIFIC PARAMS ### @@ -4198,6 +4214,9 @@ async def async_queue_request( ) # do not store the original `sk-..` api key in the db data["metadata"]["headers"] = _headers data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_team_id"] = getattr( + user_api_key_dict, "team_id", None + ) data["metadata"]["endpoint"] = str(request.url) global user_temperature, user_request_timeout, user_max_tokens, user_api_base