[Feature] image cost: token-based fallback when (quality, size) lookup misses

This commit is contained in:
kimsehwan96 2026-04-30 19:10:11 +09:00
parent d3891e6eae
commit 4ad6d201b1
3 changed files with 315 additions and 18 deletions

View file

@ -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(

View file

@ -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

View file

@ -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"])