From 8b144451e69535e0d27ded48ed54151fa1c3ab25 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 8 Dec 2023 22:00:36 -0800 Subject: [PATCH] (test) model_info calback --- litellm/tests/test_amazing_proxy_custom_logger.py | 4 ++-- litellm/tests/test_configs/test_custom_logger.yaml | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_amazing_proxy_custom_logger.py b/litellm/tests/test_amazing_proxy_custom_logger.py index 3e3e7468f90..0c4463cc7a6 100644 --- a/litellm/tests/test_amazing_proxy_custom_logger.py +++ b/litellm/tests/test_amazing_proxy_custom_logger.py @@ -78,7 +78,7 @@ def test_chat_completion(client): config_model_info = litellm_params.get("model_info") proxy_server_request_object = litellm_params.get("proxy_server_request") - assert config_model_info == {'mode': 'chat', 'input_cost_per_token': 0.0002} + assert config_model_info == {'id': 'gm', 'input_cost_per_token': 0.0002, 'mode': 'chat'} assert proxy_server_request_object == {'url': 'http://testserver/chat/completions', 'method': 'POST', 'headers': {'host': 'testserver', 'accept': '*/*', 'accept-encoding': 'gzip, deflate', 'connection': 'keep-alive', 'user-agent': 'testclient', 'authorization': 'Bearer sk-1234', 'content-length': '105', 'content-type': 'application/json'}, 'body': {'model': 'Azure OpenAI GPT-4 Canada', 'messages': [{'role': 'user', 'content': 'hi'}], 'max_tokens': 10}} result = response.json() print(f"Received response: {result}") @@ -177,7 +177,7 @@ def test_embedding(client): 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']}} - assert model_info == {'input_cost_per_token': 0.002, 'mode': 'embedding'} + assert model_info == {'input_cost_per_token': 0.002, 'mode': 'embedding', 'id': 'hello'} result = response.json() print(f"Received response: {result}") except Exception as e: diff --git a/litellm/tests/test_configs/test_custom_logger.yaml b/litellm/tests/test_configs/test_custom_logger.yaml index 4c8f414dcd5..145c618eddc 100644 --- a/litellm/tests/test_configs/test_custom_logger.yaml +++ b/litellm/tests/test_configs/test_custom_logger.yaml @@ -8,6 +8,7 @@ model_list: model_info: mode: chat input_cost_per_token: 0.0002 + id: gm - model_name: azure-embedding-model litellm_params: model: azure/azure-embedding-model @@ -17,6 +18,7 @@ model_list: model_info: mode: embedding input_cost_per_token: 0.002 + id: hello litellm_settings: drop_params: True