From 26c19ba3e162278d0ae1a7f1b8f3f2d994f300a6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 15:47:27 -0800 Subject: [PATCH 01/12] add BaseImageGenTest --- .../base_image_generation_test.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 tests/image_gen_tests/base_image_generation_test.py diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py new file mode 100644 index 000000000..56c4557cc --- /dev/null +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -0,0 +1,61 @@ +import asyncio +import httpx +import json +import pytest +import sys +from typing import Any, Dict, List +from unittest.mock import MagicMock, Mock, patch +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm +from litellm.exceptions import BadRequestError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.utils import CustomStreamWrapper +from openai.types.image import Image + +# test_example.py +from abc import ABC, abstractmethod + + +class BaseImageGenTest(ABC): + """ + Abstract base test class that enforces a common test across all test classes. + """ + + @abstractmethod + def get_base_image_generation_call_args(self) -> dict: + """Must return the base image generation call args""" + pass + + @pytest.mark.asyncio(scope="module") + async def test_basic_image_generation(self): + """Test basic image generation""" + try: + base_image_generation_call_args = self.get_base_image_generation_call_args() + litellm.set_verbose = True + response = await litellm.aimage_generation( + **base_image_generation_call_args, prompt="A image of a otter" + ) + print(response) + + assert response._hidden_params["response_cost"] is not None + from openai.types.images_response import ImagesResponse + + ImagesResponse.model_validate(response.model_dump()) + + for d in response.data: + assert isinstance(d, Image) + print("data in response.data", d) + assert d.b64_json is not None or d.url is not None + except litellm.RateLimitError as e: + pass + except litellm.ContentPolicyViolationError: + pass # Azure randomly raises these errors - skip when they occur + except Exception as e: + if "Your task failed as a result of our safety system." in str(e): + pass + else: + pytest.fail(f"An exception occurred - {str(e)}") From c15359911a87930e517077611ac2b2eed7fefa9f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 15:48:05 -0800 Subject: [PATCH 02/12] use 1 class for unit testing --- .../image_gen_tests/test_image_generation.py | 116 +++--------------- 1 file changed, 20 insertions(+), 96 deletions(-) diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index e94d62c1f..4e02b493c 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -22,6 +22,7 @@ import pytest import litellm import json import tempfile +from base_image_generation_test import BaseImageGenTest def get_vertex_ai_creds_json() -> dict: @@ -97,25 +98,27 @@ def load_vertex_ai_credentials(): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) -def test_image_generation_openai(): - try: - litellm.set_verbose = True - response = litellm.image_generation( - prompt="A cute baby sea otter", model="dall-e-3" - ) - print(f"response: {response}") - assert len(response.data) > 0 - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # OpenAI randomly raises these errors - skip when they occur - except Exception as e: - if "Connection error" in str(e): - pass - pytest.fail(f"An exception occurred - {str(e)}") +# class TestBedrockSd3(BaseImageGenTest): +# def get_base_image_generation_call_args(self) -> dict: +# return {"model": "bedrock/stability.sd3-large-v1:0"} + +# class TestBedrockSd1(BaseImageGenTest): +# def get_base_image_generation_call_args(self) -> dict: +# return {"model": "bedrock/stability.sd3-large-v1:0"} -# test_image_generation_openai() +class TestOpenAIDalle3(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + return {"model": "dall-e-3"} + + +class TestAzureOpenAIDalle3(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + return { + "model": "azure/dall-e-3-test", + "api_version": "2023-09-01-preview", + "base_model": "dall-e-3", + } @pytest.mark.parametrize( @@ -188,88 +191,9 @@ def test_image_generation_azure_dall_e_3(): pytest.fail(f"An exception occurred - {str(e)}") -# test_image_generation_azure_dall_e_3() -@pytest.mark.asyncio -async def test_async_image_generation_openai(): - try: - response = litellm.image_generation( - prompt="A cute baby sea otter", model="dall-e-3" - ) - print(f"response: {response}") - assert len(response.data) > 0 - except litellm.APIError: - pass - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # openai randomly raises these errors - skip when they occur - except litellm.InternalServerError: - pass - except Exception as e: - if "Connection error" in str(e): - pass - pytest.fail(f"An exception occurred - {str(e)}") - - # asyncio.run(test_async_image_generation_openai()) -@pytest.mark.asyncio -async def test_async_image_generation_azure(): - try: - response = await litellm.aimage_generation( - prompt="A cute baby sea otter", - model="azure/dall-e-3-test", - api_version="2023-09-01-preview", - ) - print(f"response: {response}") - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except litellm.InternalServerError: - pass - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - if "Connection error" in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", - ["bedrock/stability.sd3-large-v1:0", "bedrock/stability.stable-diffusion-xl-v1"], -) -def test_image_generation_bedrock(model): - try: - litellm.set_verbose = True - response = litellm.image_generation( - prompt="A cute baby sea otter", - model=model, - aws_region_name="us-west-2", - ) - - print(f"response: {response}") - print("response hidden params", response._hidden_params) - - assert response._hidden_params["response_cost"] is not None - from openai.types.images_response import ImagesResponse - - ImagesResponse.model_validate(response.model_dump()) - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") - - @pytest.mark.asyncio async def test_aimage_generation_bedrock_with_optional_params(): try: From a6ce4254f02690a38b475e609c634c31b5f2ed6e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:18:22 -0800 Subject: [PATCH 03/12] add debugging to BaseImageGenTest --- tests/image_gen_tests/base_image_generation_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py index 56c4557cc..be0ecb1b4 100644 --- a/tests/image_gen_tests/base_image_generation_test.py +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -42,6 +42,7 @@ class BaseImageGenTest(ABC): print(response) assert response._hidden_params["response_cost"] is not None + print("response_cost", response._hidden_params["response_cost"]) from openai.types.images_response import ImagesResponse ImagesResponse.model_validate(response.model_dump()) From 2415e7d0a6bbb95f9418e82e2df544d7d5b6ece3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:18:49 -0800 Subject: [PATCH 04/12] TestAzureOpenAIDalle3 --- tests/image_gen_tests/test_image_generation.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 4e02b493c..411aab475 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -23,6 +23,10 @@ import litellm import json import tempfile from base_image_generation_test import BaseImageGenTest +import logging +from litellm._logging import verbose_logger + +verbose_logger.setLevel(logging.DEBUG) def get_vertex_ai_creds_json() -> dict: @@ -114,10 +118,15 @@ class TestOpenAIDalle3(BaseImageGenTest): class TestAzureOpenAIDalle3(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: + litellm.set_verbose = True return { "model": "azure/dall-e-3-test", "api_version": "2023-09-01-preview", - "base_model": "dall-e-3", + "metadata": { + "model_info": { + "base_model": "dall-e-3", + } + }, } From 1c2f7f792be8725e1f5a738d1c44b1822d6db461 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:19:43 -0800 Subject: [PATCH 05/12] fix response_cost_calculator --- litellm/cost_calculator.py | 37 ++++++++++++++++++------------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 2aff3b04c..5c2884d0c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -515,6 +515,7 @@ def _infer_call_type( def completion_cost( # noqa: PLR0915 completion_response=None, model: Optional[str] = None, + base_model: Optional[str] = None, prompt="", messages: List = [], completion="", @@ -528,6 +529,7 @@ def completion_cost( # noqa: PLR0915 quality=None, n=None, # number of images ### CUSTOM PRICING ### + custom_pricing: Optional[bool] = None, custom_cost_per_token: Optional[CostPerToken] = None, custom_cost_per_second: Optional[float] = None, optional_params: Optional[dict] = None, @@ -620,7 +622,10 @@ def completion_cost( # noqa: PLR0915 f"completion_response response ms: {getattr(completion_response, '_response_ms', None)} " ) model = _select_model_name_for_cost_calc( - model=model, completion_response=completion_response + model=model, + completion_response=completion_response, + base_model=base_model, + custom_pricing=custom_pricing, ) hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: @@ -852,24 +857,18 @@ def response_cost_calculator( else: if isinstance(response_object, BaseModel): response_object._hidden_params["optional_params"] = optional_params - if isinstance(response_object, ImageResponse): - response_cost = completion_cost( - completion_response=response_object, - model=model, - call_type=call_type, - custom_llm_provider=custom_llm_provider, - optional_params=optional_params, - ) - else: - if custom_pricing is True: # override defaults if custom pricing is set - base_model = model - # base_model defaults to None if not set on model_info - response_cost = completion_cost( - completion_response=response_object, - call_type=call_type, - model=base_model, - custom_llm_provider=custom_llm_provider, - ) + + if custom_pricing is True: # override defaults if custom pricing is set + base_model = model + # base_model defaults to None if not set on model_info + response_cost = completion_cost( + completion_response=response_object, + call_type=call_type, + model=model, + base_model=base_model, + custom_llm_provider=custom_llm_provider, + custom_pricing=custom_pricing, + ) return response_cost except Exception as e: raise e From d2a1ac804b3eaf38851c3ce7cc48c153c3f5e819 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:34:47 -0800 Subject: [PATCH 06/12] test_basic_image_generation --- .../base_image_generation_test.py | 27 ++++- .../image_gen_tests/test_image_generation.py | 113 ++++-------------- 2 files changed, 47 insertions(+), 93 deletions(-) diff --git a/tests/image_gen_tests/base_image_generation_test.py b/tests/image_gen_tests/base_image_generation_test.py index be0ecb1b4..e0652114d 100644 --- a/tests/image_gen_tests/base_image_generation_test.py +++ b/tests/image_gen_tests/base_image_generation_test.py @@ -3,7 +3,7 @@ import httpx import json import pytest import sys -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from unittest.mock import MagicMock, Mock, patch import os @@ -15,6 +15,19 @@ from litellm.exceptions import BadRequestError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper from openai.types.image import Image +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardLoggingPayload + + +class TestCustomLogger(CustomLogger): + def __init__(self): + super().__init__() + self.standard_logging_payload: Optional[StandardLoggingPayload] = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_payload = kwargs.get("standard_logging_object") + pass + # test_example.py from abc import ABC, abstractmethod @@ -34,6 +47,8 @@ class BaseImageGenTest(ABC): async def test_basic_image_generation(self): """Test basic image generation""" try: + custom_logger = TestCustomLogger() + litellm.callbacks = [custom_logger] base_image_generation_call_args = self.get_base_image_generation_call_args() litellm.set_verbose = True response = await litellm.aimage_generation( @@ -41,8 +56,18 @@ class BaseImageGenTest(ABC): ) print(response) + await asyncio.sleep(1) + assert response._hidden_params["response_cost"] is not None + assert response._hidden_params["response_cost"] > 0 print("response_cost", response._hidden_params["response_cost"]) + + logged_standard_logging_payload = custom_logger.standard_logging_payload + print("logged_standard_logging_payload", logged_standard_logging_payload) + assert logged_standard_logging_payload is not None + assert logged_standard_logging_payload["response_cost"] is not None + assert logged_standard_logging_payload["response_cost"] > 0 + from openai.types.images_response import ImagesResponse ImagesResponse.model_validate(response.model_dump()) diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 411aab475..14493c3f0 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -102,13 +102,28 @@ def load_vertex_ai_credentials(): os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) -# class TestBedrockSd3(BaseImageGenTest): -# def get_base_image_generation_call_args(self) -> dict: -# return {"model": "bedrock/stability.sd3-large-v1:0"} +class TestVertexImageGeneration(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + # load_vertex_ai_credentials() + litellm.in_memory_llm_clients_cache = {} + return { + "model": "vertex_ai/imagegeneration@006", + "vertex_ai_project": "adroit-crow-413218", + "vertex_ai_location": "us-central1", + "n": 1, + } -# class TestBedrockSd1(BaseImageGenTest): -# def get_base_image_generation_call_args(self) -> dict: -# return {"model": "bedrock/stability.sd3-large-v1:0"} + +class TestBedrockSd3(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + litellm.in_memory_llm_clients_cache = {} + return {"model": "bedrock/stability.sd3-large-v1:0"} + + +class TestBedrockSd1(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + litellm.in_memory_llm_clients_cache = {} + return {"model": "bedrock/stability.sd3-large-v1:0"} class TestOpenAIDalle3(BaseImageGenTest): @@ -130,48 +145,6 @@ class TestAzureOpenAIDalle3(BaseImageGenTest): } -@pytest.mark.parametrize( - "sync_mode", - [ - True, - ], # False -) # -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -async def test_image_generation_azure(sync_mode): - try: - if sync_mode: - response = litellm.image_generation( - prompt="A cute baby sea otter", - model="azure/", - api_version="2023-06-01-preview", - ) - else: - response = await litellm.aimage_generation( - prompt="A cute baby sea otter", - model="azure/", - api_version="2023-06-01-preview", - ) - print(f"response: {response}") - assert len(response.data) > 0 - except litellm.RateLimitError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except litellm.InternalServerError: - pass - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - if "Connection error" in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") - - -# test_image_generation_azure() - - @pytest.mark.flaky(retries=3, delay=1) def test_image_generation_azure_dall_e_3(): try: @@ -221,47 +194,3 @@ async def test_aimage_generation_bedrock_with_optional_params(): pass else: pytest.fail(f"An exception occurred - {str(e)}") - - -from openai.types.image import Image - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_aimage_generation_vertex_ai(sync_mode): - - litellm.set_verbose = True - - load_vertex_ai_credentials() - data = { - "prompt": "An olympic size swimming pool", - "model": "vertex_ai/imagegeneration@006", - "vertex_ai_project": "adroit-crow-413218", - "vertex_ai_location": "us-central1", - "n": 1, - } - try: - if sync_mode: - response = litellm.image_generation(**data) - else: - response = await litellm.aimage_generation(**data) - assert response.data is not None - assert len(response.data) > 0 - - for d in response.data: - assert isinstance(d, Image) - print("data in response.data", d) - assert d.b64_json is not None - except litellm.ServiceUnavailableError as e: - pass - except litellm.RateLimitError as e: - pass - except litellm.InternalServerError as e: - pass - except litellm.ContentPolicyViolationError: - pass # Azure randomly raises these errors - skip when they occur - except Exception as e: - if "Your task failed as a result of our safety system." in str(e): - pass - else: - pytest.fail(f"An exception occurred - {str(e)}") From e0f1f18339510222ed2a8515d59b6013778ee72f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:46:21 -0800 Subject: [PATCH 07/12] fix img gen basic test --- tests/image_gen_tests/test_image_generation.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 14493c3f0..a9e5fe692 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -104,7 +104,9 @@ def load_vertex_ai_credentials(): class TestVertexImageGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: - # load_vertex_ai_credentials() + # comment this when running locally + load_vertex_ai_credentials() + litellm.in_memory_llm_clients_cache = {} return { "model": "vertex_ai/imagegeneration@006", From 501ced09ba924755c16bd7a989ccd90b199346bd Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:53:36 -0800 Subject: [PATCH 08/12] fix _select_model_name_for_cost_calc --- litellm/cost_calculator.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 5c2884d0c..11cd58c5b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -454,13 +454,14 @@ def _select_model_name_for_cost_calc( if base_model is not None: return base_model - return_model = model if isinstance(completion_response, str): return return_model elif return_model is None and hasattr(completion_response, "get"): return_model = completion_response.get("model", "") # type: ignore + elif hasattr(completion_response, "model"): + return_model = completion_response.model # type: ignore hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: @@ -857,9 +858,6 @@ def response_cost_calculator( else: if isinstance(response_object, BaseModel): response_object._hidden_params["optional_params"] = optional_params - - if custom_pricing is True: # override defaults if custom pricing is set - base_model = model # base_model defaults to None if not set on model_info response_cost = completion_cost( completion_response=response_object, From 84603763ccd32808df8b3438cc4a6c64d6853a91 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 16:58:07 -0800 Subject: [PATCH 09/12] fix test_aimage_generation_bedrock_with_optional_params --- tests/image_gen_tests/test_image_generation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index a9e5fe692..692a0e4e9 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -181,6 +181,7 @@ def test_image_generation_azure_dall_e_3(): @pytest.mark.asyncio async def test_aimage_generation_bedrock_with_optional_params(): try: + litellm.in_memory_llm_clients_cache = {} response = await litellm.aimage_generation( prompt="A cute baby sea otter", model="bedrock/stability.stable-diffusion-xl-v1", From 0a73fa6d013a5863844014b41db103676a7de2a7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 17:45:12 -0800 Subject: [PATCH 10/12] fix undo changes cost tracking --- litellm/cost_calculator.py | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 11cd58c5b..03a86fb13 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -171,7 +171,6 @@ def cost_per_token( # noqa: PLR0915 model_with_provider = model_with_provider_and_region else: _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) - model_without_prefix = model model_parts = model.split("/", 1) if len(model_parts) > 1: @@ -516,7 +515,6 @@ def _infer_call_type( def completion_cost( # noqa: PLR0915 completion_response=None, model: Optional[str] = None, - base_model: Optional[str] = None, prompt="", messages: List = [], completion="", @@ -530,7 +528,6 @@ def completion_cost( # noqa: PLR0915 quality=None, n=None, # number of images ### CUSTOM PRICING ### - custom_pricing: Optional[bool] = None, custom_cost_per_token: Optional[CostPerToken] = None, custom_cost_per_second: Optional[float] = None, optional_params: Optional[dict] = None, @@ -625,8 +622,6 @@ def completion_cost( # noqa: PLR0915 model = _select_model_name_for_cost_calc( model=model, completion_response=completion_response, - base_model=base_model, - custom_pricing=custom_pricing, ) hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: @@ -858,15 +853,24 @@ def response_cost_calculator( else: if isinstance(response_object, BaseModel): response_object._hidden_params["optional_params"] = optional_params - # base_model defaults to None if not set on model_info - response_cost = completion_cost( - completion_response=response_object, - call_type=call_type, - model=model, - base_model=base_model, - custom_llm_provider=custom_llm_provider, - custom_pricing=custom_pricing, - ) + if isinstance(response_object, ImageResponse): + response_cost = completion_cost( + completion_response=response_object, + model=model, + call_type=call_type, + custom_llm_provider=custom_llm_provider, + optional_params=optional_params, + ) + else: + if custom_pricing is True: # override defaults if custom pricing is set + base_model = model + # base_model defaults to None if not set on model_info + response_cost = completion_cost( + completion_response=response_object, + call_type=call_type, + model=base_model, + custom_llm_provider=custom_llm_provider, + ) return response_cost except Exception as e: raise e From 96ec2a5a191e307340d8c9c26317599b055de1df Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 17:46:27 -0800 Subject: [PATCH 11/12] fix response_cost_calculator --- litellm/cost_calculator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 03a86fb13..6ebec8529 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -854,6 +854,8 @@ def response_cost_calculator( if isinstance(response_object, BaseModel): response_object._hidden_params["optional_params"] = optional_params if isinstance(response_object, ImageResponse): + if base_model is not None: + model = base_model response_cost = completion_cost( completion_response=response_object, model=model, From 855761d041ee5c9f276e0170b9d8374824c8b3b6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 12 Nov 2024 17:47:31 -0800 Subject: [PATCH 12/12] fix test_cost_azure_gpt_35 --- litellm/cost_calculator.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6ebec8529..0aa8a8e36 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -459,8 +459,6 @@ def _select_model_name_for_cost_calc( elif return_model is None and hasattr(completion_response, "get"): return_model = completion_response.get("model", "") # type: ignore - elif hasattr(completion_response, "model"): - return_model = completion_response.model # type: ignore hidden_params = getattr(completion_response, "_hidden_params", None) if hidden_params is not None: