diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 33bf546c239..f4a84248c60 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1137,6 +1137,15 @@ class CostCalculatorUtils: model=model, image_response=completion_response, ) + elif custom_llm_provider == litellm.LlmProviders.TOGETHER_AI.value: + from litellm.llms.together_ai.image_generation.cost_calculator import ( + cost_calculator as together_ai_image_cost_calculator, + ) + + return together_ai_image_cost_calculator( + model=model, + image_response=completion_response, + ) elif custom_llm_provider == litellm.LlmProviders.COMETAPI.value: from litellm.llms.cometapi.image_generation.cost_calculator import ( cost_calculator as cometapi_image_cost_calculator, diff --git a/litellm/llms/together_ai/image_generation/cost_calculator.py b/litellm/llms/together_ai/image_generation/cost_calculator.py new file mode 100644 index 00000000000..406459e0e13 --- /dev/null +++ b/litellm/llms/together_ai/image_generation/cost_calculator.py @@ -0,0 +1,12 @@ +import litellm +from litellm.types.utils import ImageResponse + + +def cost_calculator(model: str, image_response: ImageResponse) -> float: + model_info = litellm.get_model_info( + model=model, + custom_llm_provider=litellm.LlmProviders.TOGETHER_AI.value, + ) + output_cost_per_image = model_info.get("output_cost_per_image") or 0.0 + num_images = len(image_response.data) if image_response.data else 0 + return output_cost_per_image * num_images diff --git a/tests/test_litellm/llms/together_ai/image_generation/test_together_ai_image_generation_cost_calculator.py b/tests/test_litellm/llms/together_ai/image_generation/test_together_ai_image_generation_cost_calculator.py new file mode 100644 index 00000000000..0fc3b2cb33b --- /dev/null +++ b/tests/test_litellm/llms/together_ai/image_generation/test_together_ai_image_generation_cost_calculator.py @@ -0,0 +1,42 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + +import litellm + +litellm.model_cost = litellm.get_model_cost_map(url="") +from litellm.llms.together_ai.image_generation.cost_calculator import cost_calculator +from litellm.types.utils import ImageObject, ImageResponse + + +def test_seedream_pricing_registered(): + info = litellm.get_model_info( + model="ByteDance-Seed/Seedream-4.0", + custom_llm_provider=litellm.LlmProviders.TOGETHER_AI.value, + ) + assert info["output_cost_per_image"] == 0.03 + assert info["mode"] == "image_generation" + + +def test_cost_calculator_scales_with_image_count(): + image_response = ImageResponse( + data=[ImageObject(url="https://x/1.png"), ImageObject(url="https://x/2.png")] + ) + cost = cost_calculator(model="ByteDance-Seed/Seedream-4.0", image_response=image_response) + assert cost == pytest.approx(0.06) + + +@pytest.mark.parametrize("call_type", ["image_generation", "image_edit"]) +def test_completion_cost_routes_together_ai_image_calls(call_type): + image_response = ImageResponse(data=[ImageObject(url="https://x/1.png")]) + cost = litellm.completion_cost( + completion_response=image_response, + model="together_ai/ByteDance-Seed/Seedream-4.0", + call_type=call_type, + ) + assert cost == pytest.approx(0.03)