From 00dfb5918c03a50091c5bc2bbcb336b55031ed1e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 16:42:49 -0800 Subject: [PATCH 1/5] (feat) proxy - log key metadata in calback --- litellm/proxy/proxy_server.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1e11e8615b..61cb74b150 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1425,6 +1425,7 @@ async def completion( data["model"] = user_model if "metadata" in data: 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"]["headers"] = dict(request.headers) else: @@ -1432,6 +1433,7 @@ async def completion( "user_api_key": user_api_key_dict.api_key, "user_api_key_user_id": user_api_key_dict.user_id, } + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) # override with user settings, these are params passed via cli if user_temperature: @@ -1588,11 +1590,13 @@ async def chat_completion( verbose_proxy_logger.debug(f'received metadata: {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_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) else: data["metadata"] = {"user_api_key": user_api_key_dict.api_key} data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata global user_temperature, user_request_timeout, user_max_tokens, user_api_base # override with user settings, these are params passed via cli @@ -1756,10 +1760,12 @@ async def embeddings( data["model"] = user_model if "metadata" in data: data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id else: data["metadata"] = {"user_api_key": user_api_key_dict.api_key} + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id @@ -1897,10 +1903,12 @@ async def image_generation( data["model"] = user_model if "metadata" in data: data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id else: data["metadata"] = {"user_api_key": user_api_key_dict.api_key} + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id @@ -2474,10 +2482,12 @@ async def async_queue_request( if "metadata" in data: verbose_proxy_logger.debug(f'received metadata: {data["metadata"]}') data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id else: data["metadata"] = {"user_api_key": user_api_key_dict.api_key} + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers) data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id From 46f84bec6938b909fa36ef7185da8a31609b2473 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 16:48:02 -0800 Subject: [PATCH 2/5] (test) api_key metadata available in callback --- litellm/tests/test_proxy_custom_logger.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/tests/test_proxy_custom_logger.py b/litellm/tests/test_proxy_custom_logger.py index e47351a9b0..34e427ef4d 100644 --- a/litellm/tests/test_proxy_custom_logger.py +++ b/litellm/tests/test_proxy_custom_logger.py @@ -182,6 +182,7 @@ def test_chat_completion(client): print("\n\n Metadata in custom logger kwargs", litellm_params.get("metadata")) assert metadata is not None assert "user_api_key" in metadata + assert "user_api_key_metadata" in metadata assert "headers" in metadata config_model_info = litellm_params.get("model_info") proxy_server_request_object = litellm_params.get("proxy_server_request") From 5c1ae3d412e9bca2bccfae7baf07c0f5240ea640 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 17:29:46 -0800 Subject: [PATCH 3/5] (feat) langfuse send metadata as tags --- litellm/integrations/langfuse.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py index db410d986f..4bf60adc94 100644 --- a/litellm/integrations/langfuse.py +++ b/litellm/integrations/langfuse.py @@ -166,13 +166,24 @@ class LangFuseLogger: input, response_obj, ): - trace = self.Langfuse.trace( - name=metadata.get("generation_name", "litellm-completion"), - input=input, - output=output, - user_id=metadata.get("trace_user_id", user_id), - id=metadata.get("trace_id", None), - ) + import langfuse + + tags = [] + supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3") + + trace_params = { + "name": metadata.get("generation_name", "litellm-completion"), + "input": input, + "output": output, + "user_id": metadata.get("trace_user_id", user_id), + "id": metadata.get("trace_id", None), + } + if supports_tags: + for key, value in metadata.items(): + tags.append(f"{key}:{value}") + trace_params.update({"tags": tags}) + + trace = self.Langfuse.trace(**trace_params) trace.generation( name=metadata.get("generation_name", "litellm-completion"), From 8df3a861784e41429e90862b03eec1abab62b3cf Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 17:44:28 -0800 Subject: [PATCH 4/5] (feat) proxy - set endpoint called in callback --- litellm/proxy/proxy_server.py | 89 ++++++++++++++--------------------- 1 file changed, 36 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61cb74b150..00acc8400b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1423,18 +1423,14 @@ async def completion( ) if user_model: data["model"] = user_model - if "metadata" in data: - 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"]["headers"] = dict(request.headers) - else: - data["metadata"] = { - "user_api_key": user_api_key_dict.api_key, - "user_api_key_user_id": user_api_key_dict.user_id, - } - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) + if "metadata" not in data: + data["metadata"] = {} + 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"]["headers"] = dict(request.headers) + data["metadata"]["endpoint"] = str(request.url) + # override with user settings, these are params passed via cli if user_temperature: data["temperature"] = user_temperature @@ -1586,17 +1582,13 @@ async def chat_completion( # if users are using user_api_key_auth, set `user` in `data` data["user"] = user_api_key_dict.user_id - if "metadata" in data: - verbose_proxy_logger.debug(f'received metadata: {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_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - else: - data["metadata"] = {"user_api_key": user_api_key_dict.api_key} - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata + if "metadata" not in data: + 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_metadata"] = user_api_key_dict.metadata + data["metadata"]["headers"] = dict(request.headers) + data["metadata"]["endpoint"] = str(request.url) global user_temperature, user_request_timeout, user_max_tokens, user_api_base # override with user settings, these are params passed via cli @@ -1758,16 +1750,13 @@ async def embeddings( ) if user_model: data["model"] = user_model - if "metadata" in data: - data["metadata"]["user_api_key"] = user_api_key_dict.api_key - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - else: - data["metadata"] = {"user_api_key": user_api_key_dict.api_key} - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + if "metadata" not in data: + data["metadata"] = {} + data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata + data["metadata"]["headers"] = dict(request.headers) + data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["endpoint"] = str(request.url) router_model_names = ( [m["model_name"] for m in llm_model_list] @@ -1901,16 +1890,14 @@ async def image_generation( ) if user_model: data["model"] = user_model - if "metadata" in data: - data["metadata"]["user_api_key"] = user_api_key_dict.api_key - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - else: + + if "metadata" not in data: data["metadata"] = {"user_api_key": user_api_key_dict.api_key} - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata + data["metadata"]["headers"] = dict(request.headers) + data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["endpoint"] = str(request.url) router_model_names = ( [m["model_name"] for m in llm_model_list] @@ -2479,17 +2466,13 @@ async def async_queue_request( # if users are using user_api_key_auth, set `user` in `data` data["user"] = user_api_key_dict.user_id - if "metadata" in data: - verbose_proxy_logger.debug(f'received metadata: {data["metadata"]}') - data["metadata"]["user_api_key"] = user_api_key_dict.api_key - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id - else: - data["metadata"] = {"user_api_key": user_api_key_dict.api_key} - data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - data["metadata"]["headers"] = dict(request.headers) - data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + if "metadata" not in data: + data["metadata"] = {} + data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata + data["metadata"]["headers"] = dict(request.headers) + data["metadata"]["user_api_key_user_id"] = user_api_key_dict.user_id + data["metadata"]["endpoint"] = str(request.url) global user_temperature, user_request_timeout, user_max_tokens, user_api_base # override with user settings, these are params passed via cli From f3a45ea044614c15577c1d9c7bd663d94031fac9 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 17 Jan 2024 17:54:18 -0800 Subject: [PATCH 5/5] (fix) cleanup --- litellm/proxy/proxy_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 00acc8400b..e014b1b6b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1892,7 +1892,7 @@ async def image_generation( data["model"] = user_model if "metadata" not in data: - data["metadata"] = {"user_api_key": user_api_key_dict.api_key} + data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata data["metadata"]["headers"] = dict(request.headers)