diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 98ca7b97ebc..d83d15cb165 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -785,6 +785,14 @@ def completion_cost( # noqa: PLR0915 raise TypeError( "completion_response must be of type ImageResponse for bedrock image cost calculation" ) + elif custom_llm_provider == litellm.LlmProviders.RECRAFT.value: + from litellm.llms.recraft.cost_calculator import ( + cost_calculator as recraft_image_cost_calculator, + ) + return recraft_image_cost_calculator( + model=model, + image_response=completion_response, + ) else: return default_image_cost_calculator( model=model, diff --git a/litellm/llms/recraft/cost_calculator.py b/litellm/llms/recraft/cost_calculator.py new file mode 100644 index 00000000000..4e57c2d9017 --- /dev/null +++ b/litellm/llms/recraft/cost_calculator.py @@ -0,0 +1,20 @@ +import litellm +from litellm.types.image import ImageResponse + + +def cost_calculator( + model: str, + image_response: ImageResponse, +) -> float: + """ + Recraft image generation cost calculator + """ + _model_info = litellm.get_model_info( + model=model, + custom_llm_provider=litellm.LlmProviders.RECRAFT.value, + ) + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 + if image_response.data: + num_images = len(image_response.data) + return output_cost_per_image * num_images