Fix: price image input tokens correctly when input_tokens_details is a dict

On the /v1/images/edits path the usage object can carry
input_tokens_details as a plain dict, and getattr() on a dict returns
None — so the text/image input split was dropped and ALL input tokens
were priced at the text rate instead of input_cost_per_image_token.
Use the existing dict-safe _get_token_detail_value helper, matching how
the output token details are already read a few lines below.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JHGpqmg9knnWYUhAD53GXo
This commit is contained in:
Ramon Emiliani 2026-08-06 16:40:16 -05:00
parent ead62528e6
commit db36bd2c8d
No known key found for this signature in database
2 changed files with 34 additions and 2 deletions

View file

@ -979,8 +979,8 @@ def calculate_image_response_cost_from_usage(
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
if input_tokens_details is not None:
prompt_tokens_details = PromptTokensDetailsWrapper(
text_tokens=getattr(input_tokens_details, "text_tokens", None),
image_tokens=getattr(input_tokens_details, "image_tokens", None),
text_tokens=_get_token_detail_value(input_tokens_details, "text_tokens"),
image_tokens=_get_token_detail_value(input_tokens_details, "image_tokens"),
cached_tokens=0,
)

View file

@ -469,5 +469,37 @@ class TestGPTImage2OutputImageTokensNoBreakdown:
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
class TestGPTImage2DictInputTokensDetails:
def test_gpt_image_2_dict_input_tokens_details_priced_at_image_rate(self):
from litellm.llms.openai.image_generation.cost_calculator import (
cost_calculator,
)
usage = Usage(
prompt_tokens=0,
completion_tokens=0,
total_tokens=3363,
input_tokens=1607,
output_tokens=1756,
input_tokens_details={"text_tokens": 163, "image_tokens": 1444},
)
image_response = ImageResponse(
created=1234567890,
data=[ImageObject(b64_json="test")],
)
image_response.usage = usage
image_response._hidden_params = {"custom_llm_provider": "openai"}
cost = cost_calculator(
model="gpt-image-2",
image_response=image_response,
custom_llm_provider="openai",
)
expected_cost = 163 * 5e-6 + 1444 * 8e-6 + 1756 * 3e-5
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
if __name__ == "__main__":
pytest.main([__file__, "-v"])