phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(feat) Add cost tracking for Azure Dall-e-3 Image Generation + use base class to ensure basic image generation tests pass (#6716)

Browse source 
* add BaseImageGenTest

* use 1 class for unit testing

* add debugging to BaseImageGenTest

* TestAzureOpenAIDalle3

* fix response_cost_calculator

* test_basic_image_generation

* fix img gen basic test

* fix _select_model_name_for_cost_calc

* fix test_aimage_generation_bedrock_with_optional_params

* fix undo changes cost tracking

* fix response_cost_calculator

* fix test_cost_azure_gpt_35
This commit is contained in:
Ishaan Jaff 2024-11-12 20:02:16 -08:00 • committed by GitHub
parent 6d4cf2d908
commit 73c7b73aa0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 139 additions and 186 deletions
Download patch file Download diff file Expand all files Collapse all files

7
litellm/cost_calculator.py
View file

@ -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:
@ -454,7 +453,6 @@ 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
@ -620,7 +618,8 @@ 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,
)
hidden_params = getattr(completion_response, "_hidden_params", None)
if hidden_params is not None:
@ -853,6 +852,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,