fix(cost): only let a WxH request size drive image pricing

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"

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