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

(feat) add cost tracking stable diffusion 3 on Bedrock (#6676)

Browse source 
* add cost tracking for sd3

* test_image_generation_bedrock

* fix get model info for image cost

* add cost_calculator for stability 1 models

* add unit testing for bedrock image cost calc

* test_cost_calculator_with_no_optional_params

* add test_cost_calculator_basic

* correctly allow size Optional

* fix cost_calculator

* sd3 unit tests cost calc
This commit is contained in:
Ishaan Jaff 2024-11-11 20:21:44 -08:00 • committed by GitHub
parent e5051a93a8
commit 25bae4cc23
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 146 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

25
litellm/cost_calculator.py
View file

@ -28,6 +28,9 @@ from litellm.llms.azure_ai.cost_calculator import (
from litellm.llms.AzureOpenAI.cost_calculation import (
cost_per_token as azure_openai_cost_per_token,
)
from litellm.llms.bedrock.image.cost_calculator import (
cost_calculator as bedrock_image_cost_calculator,
)
from litellm.llms.cohere.cost_calculator import (
cost_per_query as cohere_rerank_cost_per_query,
)
@ -521,12 +524,13 @@ def completion_cost( # noqa: PLR0915
custom_llm_provider=None,
region_name=None, # used for bedrock pricing
### IMAGE GEN ###
size=None,
size: Optional[str] = None,
quality=None,
n=None, # number of images
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
custom_cost_per_second: Optional[float] = None,
optional_params: Optional[dict] = None,
) -> float:
"""
Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm.
@ -667,7 +671,17 @@ def completion_cost( # noqa: PLR0915
# https://cloud.google.com/vertex-ai/generative-ai/pricing
# Vertex Charges Flat $0.20 per image
return 0.020
elif custom_llm_provider == "bedrock":
if isinstance(completion_response, ImageResponse):
return bedrock_image_cost_calculator(
model=model,
size=size,
image_response=completion_response,
optional_params=optional_params,
)
raise TypeError(
"completion_response must be of type ImageResponse for bedrock image cost calculation"
)
if size is None:
size = "1024-x-1024" # openai default
# fix size to match naming convention
@ -677,9 +691,9 @@ def completion_cost( # noqa: PLR0915
image_gen_model_name_with_quality = image_gen_model_name
if quality is not None:
image_gen_model_name_with_quality = f"{quality}/{image_gen_model_name}"
size = size.split("-x-")
height = int(size[0]) # if it's 1024-x-1024 vs. 1024x1024
width = int(size[1])
size_parts = size.split("-x-")
height = int(size_parts[0]) # if it's 1024-x-1024 vs. 1024x1024
width = int(size_parts[1])
verbose_logger.debug(f"image_gen_model_name: {image_gen_model_name}")
verbose_logger.debug(
f"image_gen_model_name_with_quality: {image_gen_model_name_with_quality}"
@ -844,6 +858,7 @@ def response_cost_calculator(
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