From 4ad6d201b1a57fee4c9974bcef0f5a21798cb6dc Mon Sep 17 00:00:00 2001 From: kimsehwan96 Date: Thu, 30 Apr 2026 19:10:11 +0900 Subject: [PATCH] [Feature] image cost: token-based fallback when (quality, size) lookup misses --- litellm/cost_calculator.py | 82 ++++-- .../litellm_core_utils/llm_cost_calc/utils.py | 3 + .../test_default_image_cost_calculator.py | 248 ++++++++++++++++++ 3 files changed, 315 insertions(+), 18 deletions(-) create mode 100644 tests/test_litellm/test_default_image_cost_calculator.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8a68d74be5b..d31df2a82a5 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1900,6 +1900,41 @@ def transcription_cost( ) +def _image_cost_from_token_usage( + cost_info: dict, image_response: ImageResponse +) -> Optional[float]: + """Token-based image cost from ``cost_info`` + ``image_response.usage``. + + Returns ``None`` when no token-cost keys match any non-zero token count. + """ + usage = getattr(image_response, "usage", None) + if usage is None: + return None + + def _detail(parent: str, key: str) -> int: + details = getattr(usage, parent, None) + if details is None: + return 0 + if isinstance(details, dict): + return int(details.get(key) or 0) + return int(getattr(details, key, 0) or 0) + + text_in = _detail("prompt_tokens_details", "text_tokens") + cached_in = _detail("prompt_tokens_details", "cached_tokens") + image_in = _detail("prompt_tokens_details", "image_tokens") + image_out = _detail("completion_tokens_details", "image_tokens") + text_in_uncached = max(text_in - cached_in, 0) + + rates: List[Tuple[str, int]] = [ + ("input_cost_per_token", text_in_uncached), + ("cache_read_input_token_cost", cached_in), + ("input_cost_per_image_token", image_in), + ("output_cost_per_image_token", image_out), + ] + cost = sum((cost_info.get(key) or 0) * tokens for key, tokens in rates) + return cost if cost > 0 else None + + def default_image_cost_calculator( model: str, custom_llm_provider: Optional[str] = None, @@ -1907,13 +1942,15 @@ def default_image_cost_calculator( n: Optional[int] = 1, # Default to 1 image size: Optional[str] = "1024-x-1024", # OpenAI default optional_params: Optional[dict] = None, + image_response: Optional[ImageResponse] = None, ) -> float: """ Default image cost calculator for image generation Args: model (str): Model name - image_response (ImageResponse): Response from image generation + image_response (Optional[ImageResponse]): When provided, used as a + token-based fallback if the (quality, size) lookup misses. quality (Optional[str]): Image quality setting n (Optional[int]): Number of images generated size (Optional[str]): Image size (e.g. "1024x1024" or "1024-x-1024") @@ -1978,27 +2015,36 @@ def default_image_cost_calculator( if _model is not None and _model in litellm.model_cost: cost_info = litellm.model_cost[_model] break + + # Priority 1: Use per-image pricing if available (for gpt-image-1 and similar models) + if cost_info is not None and cost_info.get("input_cost_per_image") is not None: + return cost_info["input_cost_per_image"] * n + # Priority 2: Fall back to per-pixel pricing for backward compatibility + if cost_info is not None and cost_info.get("input_cost_per_pixel") is not None: + return cost_info["input_cost_per_pixel"] * height * width * n + + # Priority 3: token-based fallback (e.g. gpt-image-2 with non-standard + # sizes where the (quality, size) chain cannot match). + if image_response is not None: + fallback_entries = [cost_info] if cost_info is not None else [] + for fallback_key in (model.split("/")[-1] if "/" in model else None, model): + if fallback_key is None: + continue + entry = litellm.model_cost.get(fallback_key) + if entry is not None and entry is not cost_info: + fallback_entries.append(entry) + for entry in fallback_entries: + cost = _image_cost_from_token_usage(entry, image_response) + if cost is not None: + return cost + if cost_info is None: raise Exception( f"Model not found in cost map. Tried checking {models_to_check}" ) - - # Priority 1: Use per-image pricing if available (for gpt-image-1 and similar models) - if ( - "input_cost_per_image" in cost_info - and cost_info["input_cost_per_image"] is not None - ): - return cost_info["input_cost_per_image"] * n - # Priority 2: Fall back to per-pixel pricing for backward compatibility - elif ( - "input_cost_per_pixel" in cost_info - and cost_info["input_cost_per_pixel"] is not None - ): - return cost_info["input_cost_per_pixel"] * height * width * n - else: - raise Exception( - f"No pricing information found for model {model}. Tried checking {models_to_check}" - ) + raise Exception( + f"No pricing information found for model {model}. Tried checking {models_to_check}" + ) def default_video_cost_calculator( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 59d0465e6d4..e6978b76440 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1002,6 +1002,7 @@ class CostCalculatorUtils: n=n, size=size, optional_params=optional_params, + image_response=completion_response, ) elif custom_llm_provider == litellm.LlmProviders.AZURE.value: # gpt-image models use token-based pricing. @@ -1024,6 +1025,7 @@ class CostCalculatorUtils: n=n, size=size, optional_params=optional_params, + image_response=completion_response, ) else: return default_image_cost_calculator( @@ -1033,5 +1035,6 @@ class CostCalculatorUtils: n=n, size=size, optional_params=optional_params, + image_response=completion_response, ) return 0.0 diff --git a/tests/test_litellm/test_default_image_cost_calculator.py b/tests/test_litellm/test_default_image_cost_calculator.py new file mode 100644 index 00000000000..1cc5c0eb328 --- /dev/null +++ b/tests/test_litellm/test_default_image_cost_calculator.py @@ -0,0 +1,248 @@ +""" +Tests for ``litellm.cost_calculator.default_image_cost_calculator``, +focused on the token-based fallback path used when the (quality, size) +lookup chain cannot resolve a per-image / per-pixel entry. + +Motivation: newer image-gen models (e.g. ``gpt-image-2``) accept +"thousands of valid resolutions" but only publish per-token pricing, +so the size-aliased lookup misses on non-standard sizes. +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../..")) + +import pytest + +import litellm +from litellm.cost_calculator import default_image_cost_calculator +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + PromptTokensDetailsWrapper, + Usage, +) + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def _image_response( + text_in: int = 0, + image_in: int = 0, + image_out: int = 0, + cached_in: int = 0, +) -> ImageResponse: + usage = Usage( + prompt_tokens=text_in + image_in, + completion_tokens=image_out, + total_tokens=text_in + image_in + image_out, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=text_in, + image_tokens=image_in, + cached_tokens=cached_in, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=0, + image_tokens=image_out, + reasoning_tokens=0, + ), + ) + response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + response.usage = usage + return response + + +class TestDefaultImageCostCalculator: + def test_per_image_entry_hits_existing_path(self): + """Regression: ``high/1024-x-1024/gpt-image-1`` resolves via the + size-aliased lookup chain and returns its ``input_cost_per_image`` + unchanged — the new fallback must not interfere. + """ + cost = default_image_cost_calculator( + model="openai/gpt-image-1", + custom_llm_provider="openai", + quality="high", + n=1, + size="1024x1024", + image_response=_image_response(text_in=10, image_out=4000), + ) + expected = litellm.model_cost["high/1024-x-1024/gpt-image-1"][ + "input_cost_per_image" + ] + assert cost == expected + + def test_per_image_entry_takes_precedence_over_token_fallback(self, monkeypatch): + """Regression: when both a (quality, size) per-image entry and a + plain token-cost entry exist for the same model, the per-image + entry wins. + """ + monkeypatch.setitem( + litellm.model_cost, + "high/1024-x-1024/synthetic-image-model", + { + "input_cost_per_image": 0.5, + "litellm_provider": "openai", + "mode": "image_generation", + }, + ) + monkeypatch.setitem( + litellm.model_cost, + "synthetic-image-model", + { + "input_cost_per_token": 5e-6, + "output_cost_per_image_token": 3e-5, + "litellm_provider": "openai", + "mode": "image_generation", + }, + ) + + cost = default_image_cost_calculator( + model="openai/synthetic-image-model", + custom_llm_provider="openai", + quality="high", + n=1, + size="1024x1024", + image_response=_image_response(text_in=100, image_out=4000), + ) + assert cost == 0.5 + + def test_token_fallback_for_non_standard_size(self, monkeypatch): + """Token-based fallback triggers when only a plain ``model`` entry + with token-cost keys is registered and the (quality, size) chain + misses. + """ + monkeypatch.setitem( + litellm.model_cost, + "synthetic-token-only-model", + { + "input_cost_per_token": 5e-6, + "output_cost_per_image_token": 3e-5, + "litellm_provider": "openai", + "mode": "image_generation", + }, + ) + + cost = default_image_cost_calculator( + model="openai/synthetic-token-only-model", + custom_llm_provider="openai", + quality="low", + n=1, + size="2048x768", + image_response=_image_response(text_in=25, image_out=772), + ) + expected = 25 * 5e-6 + 772 * 3e-5 + assert abs(cost - expected) < 1e-9 + + def test_token_fallback_includes_image_input_tokens(self, monkeypatch): + """For image-edit responses both ``image_tokens`` (input) and + ``image_tokens`` (output) must contribute to cost when their + per-token rates are registered. + """ + monkeypatch.setitem( + litellm.model_cost, + "synthetic-edit-model", + { + "input_cost_per_token": 5e-6, + "input_cost_per_image_token": 8e-6, + "output_cost_per_image_token": 3e-5, + "litellm_provider": "openai", + "mode": "image_generation", + }, + ) + + cost = default_image_cost_calculator( + model="openai/synthetic-edit-model", + custom_llm_provider="openai", + quality="medium", + n=1, + size="1280x720", + image_response=_image_response(text_in=510, image_in=1452, image_out=5488), + ) + expected = 510 * 5e-6 + 1452 * 8e-6 + 5488 * 3e-5 + assert abs(cost - expected) < 1e-9 + + def test_token_fallback_subtracts_cached_input_tokens(self, monkeypatch): + """Cached input tokens are billed at ``cache_read_input_token_cost``; + the uncached remainder uses the standard ``input_cost_per_token`` + rate so caching is reflected in the final cost. + """ + monkeypatch.setitem( + litellm.model_cost, + "synthetic-cached-model", + { + "input_cost_per_token": 5e-6, + "cache_read_input_token_cost": 1.25e-6, + "output_cost_per_image_token": 3e-5, + "litellm_provider": "openai", + "mode": "image_generation", + }, + ) + + cost = default_image_cost_calculator( + model="openai/synthetic-cached-model", + custom_llm_provider="openai", + quality="low", + n=1, + size="2048x768", + image_response=_image_response(text_in=200, cached_in=160, image_out=600), + ) + # 40 uncached text + 160 cached text + 600 image_out + expected = 40 * 5e-6 + 160 * 1.25e-6 + 600 * 3e-5 + assert abs(cost - expected) < 1e-9 + + def test_unmapped_model_without_image_response_raises(self): + """Negative: cost map miss + no ``image_response`` to fall back on + — preserve the original behaviour of raising rather than silently + returning 0. + """ + with pytest.raises(Exception, match="Model not found in cost map"): + default_image_cost_calculator( + model="openai/totally-unmapped-image-model", + custom_llm_provider="openai", + quality="high", + n=1, + size="1024x1024", + ) + + def test_entry_without_pricing_keys_raises(self, monkeypatch): + """Negative: cost map entry resolves but carries no per-image, + per-pixel, or token cost keys — ``raise`` the original + ``No pricing information found`` error. + """ + monkeypatch.setitem( + litellm.model_cost, + "high/1024-x-1024/synthetic-no-pricing-model", + { + "litellm_provider": "openai", + "mode": "image_generation", + }, + ) + + with pytest.raises(Exception, match="No pricing information found"): + default_image_cost_calculator( + model="openai/synthetic-no-pricing-model", + custom_llm_provider="openai", + quality="high", + n=1, + size="1024x1024", + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])