From bb0059542964834af8ade509e863288613d21755 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 23 Nov 2023 14:20:48 -0800 Subject: [PATCH] (feat) cost: azure gpt + testing --- litellm/tests/test_get_model_cost_map.py | 28 +++++++++++++++++++++++- litellm/utils.py | 2 +- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_get_model_cost_map.py b/litellm/tests/test_get_model_cost_map.py index e3d3235c7e..adb87e765c 100644 --- a/litellm/tests/test_get_model_cost_map.py +++ b/litellm/tests/test_get_model_cost_map.py @@ -55,4 +55,30 @@ def test_cost_ft_gpt_35(): assert cost == expected_cost except Exception as e: pytest.fail(f"Cost Calc failed for ft:gpt-3.5. Expected {expected_cost}, Calculated cost {cost}") -test_cost_ft_gpt_35() +# test_cost_ft_gpt_35() + +def test_cost_azure_gpt_35(): + try: + # this tests if litellm.completion_cost can calculate cost for azure/chatgpt-deployment-2 which maps to azure/gpt-3.5-turbo + # for this test we check if passing `model` to completion_cost overrides the completion cost + from litellm import ModelResponse, Choices, Message + from litellm.utils import Usage + resp = ModelResponse( + id='chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac', + choices=[Choices(finish_reason=None, index=0, + message=Message(content=' Sure! Here is a short poem about the sky:\n\nA canvas of blue, a', role='assistant'))], + model='chatGPT-deployment-LiteLLM-isAMAZING', + usage=Usage(prompt_tokens=21, completion_tokens=17, total_tokens=38) + ) + + cost = litellm.completion_cost(completion_response=resp, model="azure/gpt-3.5-turbo") + print("\n Calculated Cost for azure/gpt-3.5-turbo", cost) + input_cost = model_cost["azure/gpt-3.5-turbo"]["input_cost_per_token"] + output_cost = model_cost["azure/gpt-3.5-turbo"]["output_cost_per_token"] + expected_cost = (input_cost*resp.usage.prompt_tokens) + (output_cost*resp.usage.completion_tokens) + print("\n Excpected cost", expected_cost) + assert cost == expected_cost + except Exception as e: + pytest.fail(f"Cost Calc failed for azure/gpt-3.5-turbo. Expected {expected_cost}, Calculated cost {cost}") +# test_cost_azure_gpt_35() + diff --git a/litellm/utils.py b/litellm/utils.py index dce870ce6f..99ac92c688 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1678,7 +1678,7 @@ def completion_cost( # get input/output tokens from completion_response prompt_tokens = completion_response['usage']['prompt_tokens'] completion_tokens = completion_response['usage']['completion_tokens'] - model = completion_response['model'] # get model from completion_response + model = model or completion_response['model'] # check if user passed an override for model, if it's none check completion_response['model'] else: prompt_tokens = token_counter(model=model, text=prompt) completion_tokens = token_counter(model=model, text=completion)