[Fix]: Add cost tracking for image edits endpoint [OpenAI, Azure]

[ishaan-jaff]

[Fix]: Add cost tracking for image edits endpoint [OpenAI, Azure]

Relevant issues

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/litellm/ directory, Adding at least 1 test is a hard requirement - see details
• I have added a screenshot of my new test passing locally
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem

Type

🐛 Bug Fix
✅ Test

Changes

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
litellm	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	May 28, 2025 0:31am

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Pull Request Overview

This PR adds cost tracking support for the image edits endpoint on both OpenAI and Azure, including new call types, cost calculator logic, and corresponding async tests.

• Introduces image_edit and aimage_edit to CallTypes
• Extends cost calculator to include image edit responses via CostCalculatorUtils
• Adds async unit tests to verify cost logging with a custom logger for OpenAI and Azure

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
tests/image_gen_tests/test_image_edits.py	Adds test_openai_image_edit_cost_tracking and test_azure_image_edit_cost_tracking along with a TestCustomLogger class
litellm/types/utils.py	Adds new enum entries for image_edit and aimage_edit and updates ImageResponse to accept extra kwargs
litellm/litellm_core_utils/llm_cost_calc/utils.py	Implements _call_type_has_image_response to include image edit types
litellm/cost_calculator.py	Uses CostCalculatorUtils for image response checks and refactors model name handling in default_image_cost_calculator

Comments suppressed due to low confidence (2)

tests/image_gen_tests/test_image_edits.py:9

• The test references patch, AsyncMock, json, and TEST_IMAGES but none of these are imported or defined. Add import json and from unittest.mock import patch, AsyncMock, and ensure TEST_IMAGES is imported or defined.

import base64

litellm/cost_calculator.py:1140

• The type hints List and Optional are used here but not imported in this module, which will cause a NameError. Please add from typing import List, Optional at the top of the file.

models_to_check: List[Optional[str]] = [