diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 1fbc82c95a2..19e3f624268 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,6 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() +import re from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType @@ -72,11 +73,19 @@ def _get_token_detail_value(details: object, key: str) -> int | None: return value if isinstance(value, int) else None +_IMAGE_SIZE_PATTERN: Final = re.compile(r"\d+(?:x|-x-)\d+") + + 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 _requested_image_size(optional_params: Mapping[str, object] | None) -> str | None: + value: Final = _requested_image_param(optional_params, "size") + return value if value is not None and _IMAGE_SIZE_PATTERN.fullmatch(value) else None + + def get_web_search_requests(server_tool_use: Any) -> int | None: """ Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value @@ -1317,7 +1326,7 @@ class CostCalculatorUtils: ) resolved_size: Final = ( - size or completion_response.size or _requested_image_param(optional_params, "size") or "1024-x-1024" + size or completion_response.size or _requested_image_size(optional_params) or "1024-x-1024" ) resolved_quality: Final = ( quality or completion_response.quality or _requested_image_param(optional_params, "quality") or "standard" 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 722c93cdb11..58fa5720636 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 @@ -3879,3 +3879,39 @@ def test_route_image_generation_cost_falls_back_to_requested_quality( ) assert cost == expected_cost + + +@pytest.mark.parametrize( + ("requested_size", "expected_cost"), + [ + ("1536x1024", 0.05), + ("1536-x-1024", 0.05), + ("auto", 0.04), + (None, 0.04), + ], +) +def test_route_image_generation_cost_falls_back_to_requested_size(monkeypatch, requested_size, 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), + "low/1536-x-1024/grok-imagine-image-2.0": tier(0.05), + }, + ) + response = ImageResponse(data=[ImageObject(url="https://example.com/image.png")]) + optional_params = {"quality": "low", **({} if requested_size is None else {"size": requested_size})} + + 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