(test) assert metadata in callbacks

This commit is contained in:
ishaan-jaff 2023-12-11 09:47:38 -08:00
parent 0e3f7ea28f
commit 4216617688
2 changed files with 6 additions and 0 deletions

View file

@ -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

View file

@ -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']}}