fix(cost): apply discounts and margins to image and video generation costs

This commit is contained in:
Devin AI 2026-07-26 22:09:37 +00:00
parent 24123269cc
commit 6b8cc6b10e
2 changed files with 181 additions and 38 deletions

View file

@ -1105,6 +1105,57 @@ def _store_cost_breakdown_in_logging_obj(
pass
def _apply_discount_and_margin(
base_cost: float,
custom_llm_provider: str | None,
litellm_logging_obj: LitellmLoggingObject | None,
prompt_tokens_cost_usd_dollar: float = 0.0,
completion_tokens_cost_usd_dollar: float = 0.0,
cost_for_built_in_tools_cost_usd_dollar: float = 0.0,
) -> float:
"""
Apply the configured provider discount and then margin to ``base_cost`` and
persist the resulting breakdown in the logging object.
Discount is applied before margin so the margin is taken on the discounted
contract price. Returns the adjusted cost.
"""
original_cost = base_cost
(
discounted_cost,
discount_percent,
discount_amount,
) = _apply_cost_discount(
base_cost=base_cost,
custom_llm_provider=custom_llm_provider,
)
(
final_cost,
margin_percent,
margin_fixed_amount,
margin_total_amount,
) = _apply_cost_margin(
base_cost=discounted_cost,
custom_llm_provider=custom_llm_provider,
)
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar,
completion_tokens_cost_usd_dollar=completion_tokens_cost_usd_dollar,
cost_for_built_in_tools_cost_usd_dollar=cost_for_built_in_tools_cost_usd_dollar,
total_cost_usd_dollar=final_cost,
original_cost=original_cost,
discount_percent=discount_percent,
discount_amount=discount_amount,
margin_percent=margin_percent,
margin_fixed_amount=margin_fixed_amount,
margin_total_amount=margin_total_amount,
)
return final_cost
def completion_cost(
completion_response=None,
model: Optional[str] = None,
@ -1335,7 +1386,7 @@ def completion_cost(
completion_response, ImageResponse
):
### IMAGE GENERATION COST CALCULATION ###
return CostCalculatorUtils.route_image_generation_cost_calculator(
image_generation_cost = CostCalculatorUtils.route_image_generation_cost_calculator(
model=model,
custom_llm_provider=custom_llm_provider,
completion_response=completion_response,
@ -1345,6 +1396,11 @@ def completion_cost(
optional_params=optional_params,
call_type=call_type,
)
return _apply_discount_and_margin(
base_cost=image_generation_cost,
custom_llm_provider=custom_llm_provider,
litellm_logging_obj=litellm_logging_obj,
)
elif call_type in _VIDEO_CALL_TYPES:
### VIDEO GENERATION COST CALCULATION ###
# Extract custom model_info for deployment-specific pricing
@ -1375,21 +1431,31 @@ def completion_cost(
video_generation_cost,
)
return video_generation_cost(
video_cost = video_generation_cost(
model=model,
duration_seconds=duration_seconds,
custom_llm_provider=custom_llm_provider,
model_info=_video_model_info,
video_resolution=video_resolution,
)
return _apply_discount_and_margin(
base_cost=video_cost,
custom_llm_provider=custom_llm_provider,
litellm_logging_obj=litellm_logging_obj,
)
# Fallback to default video cost calculation if no duration available
return default_video_cost_calculator(
video_cost = default_video_cost_calculator(
model=model,
duration_seconds=0.0, # Default to 0 if no duration available
custom_llm_provider=custom_llm_provider,
model_info=_video_model_info,
video_resolution=video_resolution,
)
return _apply_discount_and_margin(
base_cost=video_cost,
custom_llm_provider=custom_llm_provider,
litellm_logging_obj=litellm_logging_obj,
)
elif call_type in _SPEECH_CALL_TYPES:
prompt_characters = litellm.utils._count_characters(text=prompt)
elif call_type in _TRANSCRIPTION_CALL_TYPES:
@ -1446,46 +1512,13 @@ def completion_cost(
)
# Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost)
_final_cost = prompt_cost + completion_cost_result
# Apply discount
original_cost = _final_cost
(
_final_cost,
discount_percent,
discount_amount,
) = _apply_cost_discount(
base_cost=_final_cost,
return _apply_discount_and_margin(
base_cost=prompt_cost + completion_cost_result,
custom_llm_provider=custom_llm_provider,
)
# Apply margin from module-level config if configured
(
_final_cost,
margin_percent,
margin_fixed_amount,
margin_total_amount,
) = _apply_cost_margin(
base_cost=_final_cost,
custom_llm_provider=custom_llm_provider,
)
# Store cost breakdown in logging object if available
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=prompt_cost,
completion_tokens_cost_usd_dollar=completion_cost_result,
cost_for_built_in_tools_cost_usd_dollar=0.0,
total_cost_usd_dollar=_final_cost,
original_cost=original_cost,
discount_percent=discount_percent,
discount_amount=discount_amount,
margin_percent=margin_percent,
margin_fixed_amount=margin_fixed_amount,
margin_total_amount=margin_total_amount,
)
return _final_cost
elif call_type == _AREALTIME_CALL_TYPE and isinstance(
completion_response, LiteLLMRealtimeStreamLoggingObject
):

View file

@ -2135,6 +2135,116 @@ def test_cost_margin_with_discount():
print(f" - Expected: ${expected_cost:.6f}")
def test_image_generation_cost_applies_discount_and_margin():
"""
Regression test for https://github.com/BerriAI/litellm/issues/34731
Image generation cost must honor cost_discount_config and cost_margin_config
the same way chat completion cost does. The image branch previously returned
the raw provider cost before either adjustment ran.
"""
from unittest.mock import MagicMock
from litellm import completion_cost
from litellm.types.utils import ImageObject, ImageResponse
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
original_discount_config = litellm.cost_discount_config.copy()
original_margin_config = litellm.cost_margin_config.copy()
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(b64_json=None, revised_prompt=None, url="https://example.com/image.png")],
quality="standard",
size="1024x1024",
)
completion_cost_kwargs = {
"completion_response": image_response,
"model": "dall-e-3",
"custom_llm_provider": "openai",
"call_type": "image_generation",
"size": "1024x1024",
"quality": "standard",
"n": 1,
"litellm_logging_obj": MagicMock(),
}
try:
litellm.cost_discount_config = {}
litellm.cost_margin_config = {}
raw_cost = completion_cost(**completion_cost_kwargs)
litellm.cost_discount_config = {"openai": 0.10}
litellm.cost_margin_config = {"openai": 0.20}
configured_cost = completion_cost(**completion_cost_kwargs)
finally:
litellm.cost_discount_config = original_discount_config
litellm.cost_margin_config = original_margin_config
assert raw_cost > 0
# discount is applied before margin: raw * (1 - 0.10) * (1 + 0.20)
expected_cost = raw_cost * 0.90 * 1.20
assert configured_cost == pytest.approx(expected_cost, rel=1e-9)
assert configured_cost != pytest.approx(raw_cost, rel=1e-9)
def test_video_generation_cost_applies_discount_and_margin():
"""
Regression test for https://github.com/BerriAI/litellm/issues/34731
Video generation cost must honor cost_discount_config and cost_margin_config.
The video branch previously returned the raw provider cost before either
adjustment ran.
"""
from unittest.mock import MagicMock
from litellm.cost_calculator import completion_cost
original_discount_config = litellm.cost_discount_config.copy()
original_margin_config = litellm.cost_margin_config.copy()
mock_response = MagicMock()
mock_response.usage = MagicMock()
mock_response.usage.duration_seconds = 10.0
mock_response.usage.video_resolution = None
type(mock_response)._hidden_params = {}
mock_logging_obj = MagicMock()
mock_logging_obj.litellm_params = {
"metadata": {"model_info": {"output_cost_per_video_per_second": 0.05}}
}
completion_cost_kwargs = {
"completion_response": mock_response,
"model": "openai/hunyuanvideo",
"call_type": "create_video",
"custom_llm_provider": "openai",
"custom_pricing": True,
"litellm_logging_obj": mock_logging_obj,
}
try:
litellm.cost_discount_config = {}
litellm.cost_margin_config = {}
raw_cost = completion_cost(**completion_cost_kwargs)
litellm.cost_discount_config = {"openai": 0.10}
litellm.cost_margin_config = {"openai": 0.20}
configured_cost = completion_cost(**completion_cost_kwargs)
finally:
litellm.cost_discount_config = original_discount_config
litellm.cost_margin_config = original_margin_config
assert raw_cost == pytest.approx(0.5, rel=1e-9)
# discount is applied before margin: raw * (1 - 0.10) * (1 + 0.20)
expected_cost = raw_cost * 0.90 * 1.20
assert configured_cost == pytest.approx(expected_cost, rel=1e-9)
assert configured_cost != pytest.approx(raw_cost, rel=1e-9)
def test_azure_image_generation_cost_calculator():
from unittest.mock import MagicMock