From 71eee5087d4bff283206df2f7452f24deb4773ba Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 11 Dec 2023 09:47:38 -0800 Subject: [PATCH] (test) assert metadata in callbacks --- litellm/proxy/proxy_server.py | 2 ++ litellm/tests/test_amazing_proxy_custom_logger.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 259ead7c945..6bab8dc8d84 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1098,8 +1098,10 @@ async def embeddings(request: Request, user_api_key_dict: UserAPIKeyAuth = Depen data["model"] = user_model if "metadata" in data: data["metadata"]["user_api_key"] = user_api_key_dict.api_key + data["metadata"]["headers"] = dict(request.headers) else: data["metadata"] = {"user_api_key": user_api_key_dict.api_key} + data["metadata"]["headers"] = dict(request.headers) router_model_names = [m["model_name"] for m in llm_model_list] if llm_model_list is not None else [] print(f"received data: {data['input']}") if "input" in data and isinstance(data['input'], list) and isinstance(data['input'][0], list) and isinstance(data['input'][0][0], int): # check if array of tokens passed in diff --git a/litellm/tests/test_amazing_proxy_custom_logger.py b/litellm/tests/test_amazing_proxy_custom_logger.py index 14db8e2af24..3b0506cae3b 100644 --- a/litellm/tests/test_amazing_proxy_custom_logger.py +++ b/litellm/tests/test_amazing_proxy_custom_logger.py @@ -77,6 +77,8 @@ def test_chat_completion(client): metadata = litellm_params.get("metadata", None) print("\n\n Metadata in custom logger kwargs", litellm_params.get("metadata")) assert metadata is not None + assert "user_api_key" in metadata + assert "headers" in metadata config_model_info = litellm_params.get("model_info") proxy_server_request_object = litellm_params.get("proxy_server_request") @@ -179,6 +181,8 @@ def test_embedding(client): metadata = litellm_params.get("metadata", None) print("\n\n Metadata in custom logger kwargs", litellm_params.get("metadata")) assert metadata is not None + assert "user_api_key" in metadata + assert "headers" in metadata proxy_server_request = litellm_params.get("proxy_server_request") model_info = litellm_params.get("model_info") assert proxy_server_request == {'url': 'http://testserver/embeddings', 'method': 'POST', 'headers': {'host': 'testserver', 'accept': '*/*', 'accept-encoding': 'gzip, deflate', 'connection': 'keep-alive', 'user-agent': 'testclient', 'authorization': 'Bearer sk-1234', 'content-length': '54', 'content-type': 'application/json'}, 'body': {'model': 'azure-embedding-model', 'input': ['hello']}}