diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index bdbaee00c19..1fbc82c95a2 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -72,6 +72,11 @@ def _get_token_detail_value(details: object, key: str) -> int | None: return value if isinstance(value, int) else None +def _requested_image_param(optional_params: Mapping[str, object] | None, key: str) -> str | None: + value: Final = None if optional_params is None else optional_params.get(key) + return value if isinstance(value, str) else None + + def get_web_search_requests(server_tool_use: Any) -> int | None: """ Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value @@ -1311,12 +1316,13 @@ class CostCalculatorUtils: cost_calculator as vertex_ai_image_cost_calculator, ) - if size is None: - size = completion_response.size or "1024-x-1024" - if quality is None: - quality = completion_response.quality or "standard" - if n is None: - n = len(completion_response.data) if completion_response.data else 0 + resolved_size: Final = ( + size or completion_response.size or _requested_image_param(optional_params, "size") or "1024-x-1024" + ) + resolved_quality: Final = ( + quality or completion_response.quality or _requested_image_param(optional_params, "quality") or "standard" + ) + resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0) if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value: if isinstance(completion_response, ImageResponse): @@ -1328,7 +1334,7 @@ class CostCalculatorUtils: if isinstance(completion_response, ImageResponse): return bedrock_image_cost_calculator( model=model, - size=size, + size=resolved_size, image_response=completion_response, optional_params=optional_params, ) @@ -1424,19 +1430,19 @@ class CostCalculatorUtils: # Fall through to default for DALL-E models return default_image_cost_calculator( model=model, - quality=quality, + quality=resolved_quality, custom_llm_provider=custom_llm_provider, - n=n, - size=size, + n=resolved_n, + size=resolved_size, optional_params=optional_params, ) else: return default_image_cost_calculator( model=model, - quality=quality, + quality=resolved_quality, custom_llm_provider=custom_llm_provider, - n=n, - size=size, + n=resolved_n, + size=resolved_size, optional_params=optional_params, ) return 0.0 diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index ce90719789a..722c93cdb11 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -27,6 +27,7 @@ from litellm.types.utils import ( ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( + CostCalculatorUtils, PromptTokensDetailsResult, TokenTypeCostBreakdown, _calculate_input_cost, @@ -3841,3 +3842,40 @@ def test_generic_cost_per_token_grok_46_long_context(_local_model_cost_map): ) assert prompt_cost == pytest.approx(200_000 * 4e-06 + 50_000 * 1e-06) assert completion_cost == pytest.approx(1_000 * 1.2e-05) + + +@pytest.mark.parametrize( + ("response_quality", "requested_quality", "expected_cost"), + [ + (None, "low", 0.04), + (None, None, 0.06), + ("high", "low", 0.08), + ], +) +def test_route_image_generation_cost_falls_back_to_requested_quality( + monkeypatch, response_quality, requested_quality, expected_cost +): + def tier(cost): + return {"litellm_provider": "xai", "mode": "image_generation", "input_cost_per_image": cost} + + monkeypatch.setattr( + litellm, + "model_cost", + { + "xai/grok-imagine-image-2.0": tier(0.06), + "low/1024-x-1024/grok-imagine-image-2.0": tier(0.04), + "high/1024-x-1024/grok-imagine-image-2.0": tier(0.08), + }, + ) + response = ImageResponse(data=[ImageObject(url="https://example.com/image.png")], quality=response_quality) + optional_params = {} if requested_quality is None else {"quality": requested_quality} + + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="xai/grok-imagine-image-2.0", + completion_response=response, + custom_llm_provider="xai", + optional_params=optional_params, + call_type="image_generation", + ) + + assert cost == expected_cost