From cc7a690c9b7f1d1583f006c83f83e254ee846c7b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 7 Feb 2024 18:18:15 -0800 Subject: [PATCH] (fix) azure_base_model cost calc --- litellm/tests/test_custom_callback_router.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/litellm/tests/test_custom_callback_router.py b/litellm/tests/test_custom_callback_router.py index b33cfd8cf02..071d4529db7 100644 --- a/litellm/tests/test_custom_callback_router.py +++ b/litellm/tests/test_custom_callback_router.py @@ -269,14 +269,15 @@ class CompletionCustomHandler( assert isinstance(kwargs["model"], str) # checking we use base_model for azure cost calculation - base_model = ( - kwargs.get("litellm_params", {}) - .get("metadata", {}) - .get("model_info", {}) - .get("base_model", None) + base_model = litellm.utils._get_base_model_from_metadata( + model_call_details=kwargs ) - if kwargs["model"] == "chatgpt-v-2" and base_model is not None: + if ( + kwargs["model"] == "chatgpt-v-2" + and base_model is not None + and kwargs["stream"] != True + ): # when base_model is set for azure, we should use pricing for the base_model # this checks response_cost == litellm.cost_per_token(model=base_model) assert isinstance(kwargs["response_cost"], float)