refactor(chatgpt): split image transformation modules

litellm/__init__.py

@@ -1871,8 +1871,10 @@ if TYPE_CHECKING:
     from .llms.chatgpt.responses.transformation import (
         ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig,
     )
-    from .llms.chatgpt.image_generation.transformation import (
+    from .llms.chatgpt.image_edit import (
         ChatGPTImageEditConfig as ChatGPTImageEditConfig,
+    )
+    from .llms.chatgpt.image_generation import (
         ChatGPTImageGenerationConfig as ChatGPTImageGenerationConfig,
     )
     from .llms.gigachat.chat.transformation import GigaChatConfig as GigaChatConfig

litellm/_lazy_imports_registry.py

@@ -1132,11 +1132,11 @@ _LLM_CONFIGS_IMPORT_MAP = {
         "ChatGPTResponsesAPIConfig",
     ),
     "ChatGPTImageGenerationConfig": (
-        ".llms.chatgpt.image_generation.transformation",
+        ".llms.chatgpt.image_generation.generation_transformation",
         "ChatGPTImageGenerationConfig",
     ),
     "ChatGPTImageEditConfig": (
-        ".llms.chatgpt.image_generation.transformation",
+        ".llms.chatgpt.image_edit.transformation",
         "ChatGPTImageEditConfig",
     ),
     "NebiusConfig": (".llms.nebius.chat.transformation", "NebiusConfig"),

litellm/llms/chatgpt/image_edit/__init__.py

@@ -0,0 +1,3 @@
+from .transformation import ChatGPTImageEditConfig
+
+__all__ = ["ChatGPTImageEditConfig"]

litellm/llms/chatgpt/image_edit/transformation.py

@@ -0,0 +1,179 @@
+import base64
+from io import BufferedReader, BytesIO
+from os import PathLike
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import httpx
+from httpx._types import RequestFiles
+
+from litellm.images.utils import ImageEditRequestUtils
+from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+from litellm.llms.chatgpt.image_generation.generation_transformation import (
+    ChatGPTImageGenerationConfig,
+)
+from litellm.llms.openai.common_utils import OpenAIError
+from litellm.types.images.main import ImageEditOptionalRequestParams
+from litellm.types.llms.openai import FileTypes
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import ImageResponse
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+
+
+class ChatGPTImageEditConfig(BaseImageEditConfig):
+    """
+    Bridge OpenAI-style Images Edits calls to ChatGPT/Codex Responses image generation.
+    """
+
+    def __init__(self) -> None:
+        self.image_generation_config = ChatGPTImageGenerationConfig()
+
+    def get_supported_openai_params(self, model: str) -> List[str]:
+        return ["size"]
+
+    def map_openai_params(
+        self,
+        image_edit_optional_params: ImageEditOptionalRequestParams,
+        model: str,
+        drop_params: bool,
+    ) -> Dict[str, Any]:
+        supported_params = self.get_supported_openai_params(model)
+        return {
+            key: value
+            for key, value in image_edit_optional_params.items()
+            if key in supported_params
+        }
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+        litellm_params: Optional[dict] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        return self.image_generation_config.validate_environment(
+            headers=headers,
+            model=model,
+            messages=[],
+            optional_params={},
+            litellm_params=litellm_params or {},
+            api_key=api_key,
+            api_base=api_base,
+        )
+
+    def get_complete_url(
+        self,
+        model: str,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        return self.image_generation_config.get_complete_url(
+            api_base=api_base,
+            api_key=litellm_params.get("api_key"),
+            model=model,
+            optional_params={},
+            litellm_params=litellm_params,
+        )
+
+    def use_multipart_form_data(self) -> bool:
+        return False
+
+    def transform_image_edit_request(
+        self,
+        model: str,
+        prompt: Optional[str],
+        image: Optional[FileTypes],
+        image_edit_optional_request_params: Dict,
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,
+    ) -> Tuple[Dict[str, Any], RequestFiles]:
+        optional_params = dict(image_edit_optional_request_params)
+        self.image_generation_config._validate_openai_image_generation_params(
+            model, optional_params
+        )
+
+        input_images = self._prepare_input_images(image)
+        if not input_images:
+            raise ValueError("ChatGPT image edit requires at least one image.")
+
+        request = self.image_generation_config._build_responses_image_request(
+            model=model,
+            prompt=prompt,
+            optional_params=optional_params,
+            litellm_params=dict(litellm_params),
+            input_images=input_images,
+        )
+        return request, []
+
+    def transform_image_edit_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        logging_obj: "LiteLLMLoggingObj",
+    ) -> ImageResponse:
+        return self.image_generation_config.transform_image_generation_response(
+            model=model,
+            raw_response=raw_response,
+            model_response=ImageResponse(),
+            logging_obj=logging_obj,
+            request_data={},
+            optional_params={},
+            litellm_params={},
+            encoding=None,
+        )
+
+    def _prepare_input_images(
+        self, image: Optional[Union[FileTypes, List[FileTypes]]]
+    ) -> List[Dict[str, Any]]:
+        if image is None:
+            return []
+
+        images = image if isinstance(image, list) else [image]
+        input_images: List[Dict[str, Any]] = []
+        for img in images:
+            if img is None:
+                continue
+            mime_type = ImageEditRequestUtils.get_image_content_type(img)
+            image_bytes = self._read_image_bytes(img)
+            b64_data = base64.b64encode(image_bytes).decode("utf-8")
+            input_images.append(
+                {
+                    "type": "input_image",
+                    "image_url": f"data:{mime_type};base64,{b64_data}",
+                }
+            )
+        return input_images
+
+    @staticmethod
+    def _read_image_bytes(image: FileTypes) -> bytes:
+        if isinstance(image, bytes):
+            return image
+        if isinstance(image, BytesIO):
+            current_pos = image.tell()
+            image.seek(0)
+            data = image.read()
+            image.seek(current_pos)
+            return data
+        if isinstance(image, BufferedReader):
+            current_pos = image.tell()
+            image.seek(0)
+            data = image.read()
+            image.seek(current_pos)
+            return data
+        if isinstance(image, tuple):
+            return ChatGPTImageEditConfig._read_image_bytes(image[1])
+        if isinstance(image, PathLike):
+            with open(image, "rb") as image_file:
+                return image_file.read()
+        raise ValueError("Unsupported image type for ChatGPT image edit.")
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+    ) -> OpenAIError:
+        return self.image_generation_config.get_error_class(
+            error_message=error_message,
+            status_code=status_code,
+            headers=headers,
+        )

litellm/llms/chatgpt/image_generation/__init__.py

@@ -1,3 +1,3 @@
-from .transformation import ChatGPTImageEditConfig, ChatGPTImageGenerationConfig
+from .generation_transformation import ChatGPTImageGenerationConfig
 
-__all__ = ["ChatGPTImageEditConfig", "ChatGPTImageGenerationConfig"]
+__all__ = ["ChatGPTImageGenerationConfig"]

litellm/llms/chatgpt/image_generation/generation_transformation.py

@@ -0,0 +1,319 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import httpx
+
+from litellm.exceptions import AuthenticationError
+from litellm.llms.base_llm.image_generation.transformation import (
+    BaseImageGenerationConfig,
+)
+from litellm.llms.openai.common_utils import OpenAIError
+from litellm.types.llms.openai import (
+    AllMessageValues,
+    OpenAIImageGenerationOptionalParams,
+)
+from litellm.types.utils import ImageObject, ImageResponse, ImageUsage
+
+from ..authenticator import Authenticator
+from ..common_utils import (
+    CHATGPT_API_BASE,
+    GetAccessTokenError,
+    ensure_chatgpt_session_id,
+    get_chatgpt_default_headers,
+    get_chatgpt_default_instructions,
+)
+from .response_parsing import (
+    dedupe,
+    extract_image_payloads,
+    extract_image_usage,
+    extract_images_from_nested_value,
+    extract_images_from_payload,
+    get_image_generation_usage,
+    get_image_strings_from_dict,
+    get_parsed_payloads,
+    is_zero_image_usage,
+    looks_like_sse,
+    parse_sse_payloads,
+    transform_image_usage,
+)
+
+GPT_IMAGE_MODEL_PREFIX = "gpt-image-"
+
+ALLOWED_OUTPUT_FORMATS = {"png", "jpeg", "webp"}
+INTERNAL_OPTIONAL_PARAMS = {"chatgpt_responses_model"}
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+
+
+class ChatGPTImageGenerationConfig(BaseImageGenerationConfig):
+    """
+    Bridge OpenAI-style Images API calls to ChatGPT/Codex Responses image generation.
+    """
+
+    def __init__(self) -> None:
+        self.authenticator = Authenticator()
+
+    def get_supported_openai_params(
+        self, model: str
+    ) -> List[OpenAIImageGenerationOptionalParams]:
+        return [
+            "output_format",
+            "size",
+        ]
+
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        supported_params = self.get_supported_openai_params(model)
+        for key, value in non_default_params.items():
+            if key in optional_params:
+                continue
+            if key in supported_params:
+                optional_params[key] = value
+            elif drop_params:
+                continue
+            else:
+                raise ValueError(
+                    f"Parameter {key} is not supported for model {model}. "
+                    f"Supported parameters are {supported_params}. "
+                    "Set drop_params=True to drop unsupported parameters."
+                )
+        return optional_params
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        try:
+            access_token = self.authenticator.get_access_token()
+        except GetAccessTokenError as e:
+            raise AuthenticationError(
+                model=model,
+                llm_provider="chatgpt",
+                message=str(e),
+            )
+
+        account_id = self.authenticator.get_account_id()
+        session_id = ensure_chatgpt_session_id(litellm_params)
+        default_headers = get_chatgpt_default_headers(
+            access_token, account_id, session_id
+        )
+        return {**default_headers, **headers}
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        api_base = self.authenticator.get_api_base() or CHATGPT_API_BASE
+        api_base = self._canonicalize_codex_api_base(api_base)
+        return f"{api_base}/responses"
+
+    @staticmethod
+    def _canonicalize_codex_api_base(api_base: str) -> str:
+        api_base = api_base.rstrip("/")
+        if api_base.endswith("/responses"):
+            api_base = api_base[: -len("/responses")]
+        if api_base.endswith("/backend-api"):
+            return f"{api_base}/codex"
+        return api_base
+
+    def transform_image_generation_request(
+        self,
+        model: str,
+        prompt: str,
+        optional_params: dict,
+        litellm_params: dict,
+        headers: dict,
+    ) -> dict:
+        self._validate_openai_image_generation_params(model, optional_params)
+
+        return self._build_responses_image_request(
+            model=model,
+            prompt=prompt,
+            optional_params=optional_params,
+            litellm_params=litellm_params,
+        )
+
+    def _build_responses_image_request(
+        self,
+        model: str,
+        prompt: Optional[str],
+        optional_params: dict,
+        litellm_params: dict,
+        input_images: Optional[List[Dict[str, Any]]] = None,
+    ) -> dict:
+        # Intentionally pinned fallback for ChatGPT image generation through the
+        # Codex Responses API. Users can override this per request or via
+        # litellm_params.
+        responses_model = (
+            optional_params.pop("chatgpt_responses_model", None)
+            or litellm_params.get("chatgpt_responses_model")
+            or "gpt-5.5"
+        )
+        content: List[Dict[str, Any]] = []
+        if prompt:
+            content.append({"type": "input_text", "text": prompt})
+        if input_images:
+            content.extend(input_images)
+
+        request: Dict[str, Any] = {
+            "model": responses_model,
+            "input": [
+                {
+                    "role": "user",
+                    "content": content,
+                }
+            ],
+            "instructions": get_chatgpt_default_instructions(),
+            "tools": [{"type": "image_generation", "model": model}],
+            "tool_choice": {"type": "image_generation"},
+            "stream": True,
+            "store": False,
+        }
+
+        image_tool = request["tools"][0]
+        for key in (
+            "output_format",
+            "size",
+        ):
+            if optional_params.get(key) is not None:
+                image_tool[key] = optional_params[key]
+
+        return request
+
+    def _validate_openai_image_generation_params(
+        self, model: str, optional_params: dict
+    ) -> None:
+        if not model.startswith(GPT_IMAGE_MODEL_PREFIX):
+            raise ValueError(
+                "ChatGPT image generation requires a GPT Image model "
+                "(for example gpt-image-1.5 or gpt-image-2)."
+            )
+
+        supported_params = set(self.get_supported_openai_params(model))
+        unsupported_params = [
+            key
+            for key in optional_params
+            if key not in supported_params and key not in INTERNAL_OPTIONAL_PARAMS
+        ]
+        if unsupported_params:
+            raise ValueError(
+                f"Parameters {unsupported_params} are not supported for model {model}. "
+                f"Supported parameters are {sorted(supported_params)}."
+            )
+
+        output_format = optional_params.get("output_format")
+        if output_format is not None and output_format not in ALLOWED_OUTPUT_FORMATS:
+            raise ValueError("output_format must be one of png, jpeg, or webp")
+
+    def transform_image_generation_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: ImageResponse,
+        logging_obj: "LiteLLMLoggingObj",
+        request_data: dict,
+        optional_params: dict,
+        litellm_params: dict,
+        encoding: Any,
+        api_key: Optional[str] = None,
+        json_mode: Optional[bool] = None,
+    ) -> ImageResponse:
+        logging_obj.post_call(
+            input=request_data.get("input", ""),
+            api_key=api_key,
+            additional_args={"complete_input_dict": request_data},
+            original_response=raw_response.text,
+        )
+
+        image_payloads = self._extract_image_payloads(raw_response)
+        if not image_payloads:
+            raise OpenAIError(
+                message="No image data found in ChatGPT image generation response",
+                status_code=raw_response.status_code,
+            )
+
+        response = ImageResponse(
+            data=[
+                ImageObject(b64_json=image_payload) for image_payload in image_payloads
+            ]
+        )
+        response.usage = None
+        image_usage = self._extract_image_usage(raw_response)
+        if image_usage is not None:
+            response.usage = image_usage
+        response.size = optional_params.get("size")
+        response.output_format = optional_params.get("output_format")
+        response._hidden_params["model"] = model
+        return response
+
+    def _extract_image_payloads(self, raw_response: httpx.Response) -> List[str]:
+        return extract_image_payloads(raw_response)
+
+    def _extract_image_usage(
+        self, raw_response: httpx.Response
+    ) -> Optional[ImageUsage]:
+        return extract_image_usage(raw_response)
+
+    def _get_parsed_payloads(self, raw_response: httpx.Response) -> List[dict]:
+        return get_parsed_payloads(raw_response)
+
+    @staticmethod
+    def _transform_image_usage(usage: dict) -> ImageUsage:
+        return transform_image_usage(usage)
+
+    @staticmethod
+    def _get_image_generation_usage(response_payload: Any) -> Optional[dict]:
+        return get_image_generation_usage(response_payload)
+
+    @staticmethod
+    def _is_zero_image_usage(usage: dict) -> bool:
+        return is_zero_image_usage(usage)
+
+    @staticmethod
+    def _looks_like_sse(body_text: str) -> bool:
+        return looks_like_sse(body_text)
+
+    @staticmethod
+    def _parse_sse_payloads(body_text: str) -> List[dict]:
+        return parse_sse_payloads(body_text)
+
+    def _extract_images_from_payload(
+        self, payload: dict
+    ) -> Tuple[List[str], List[str]]:
+        return extract_images_from_payload(payload)
+
+    def _extract_images_from_nested_value(self, value: Any) -> List[str]:
+        return extract_images_from_nested_value(value)
+
+    @staticmethod
+    def _get_image_strings_from_dict(value: dict) -> List[str]:
+        return get_image_strings_from_dict(value)
+
+    @staticmethod
+    def _dedupe(values: List[str]) -> List[str]:
+        return dedupe(values)
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+    ) -> OpenAIError:
+        return OpenAIError(
+            message=error_message,
+            status_code=status_code,
+            headers=headers,
+        )

litellm/llms/chatgpt/image_generation/response_parsing.py

@@ -0,0 +1,235 @@
+import json
+from typing import Any, List, Optional, Tuple
+
+import httpx
+
+from litellm.constants import STREAM_SSE_DONE_STRING
+from litellm.llms.openai.common_utils import OpenAIError
+from litellm.types.llms.openai import ResponsesAPIStreamEvents
+from litellm.types.utils import ImageUsage, ImageUsageInputTokensDetails
+from litellm.utils import CustomStreamWrapper
+
+
+def extract_image_payloads(raw_response: httpx.Response) -> List[str]:
+    content_type = raw_response.headers.get("content-type", "")
+    body_text = raw_response.text or ""
+    parsed_payloads: List[dict] = []
+
+    if "text/event-stream" in content_type.lower() or looks_like_sse(body_text):
+        parsed_payloads = parse_sse_payloads(body_text)
+    else:
+        try:
+            response_json = raw_response.json()
+        except Exception:
+            response_json = {}
+        if isinstance(response_json, dict):
+            parsed_payloads = [response_json]
+
+    images: List[str] = []
+    partial_images: List[str] = []
+    for payload in parsed_payloads:
+        extracted_images, extracted_partial_images = extract_images_from_payload(
+            payload
+        )
+        images.extend(extracted_images)
+        partial_images.extend(extracted_partial_images)
+    return dedupe(images) or dedupe(partial_images)
+
+
+def extract_image_usage(raw_response: httpx.Response) -> Optional[ImageUsage]:
+    parsed_payloads = get_parsed_payloads(raw_response)
+
+    for payload in parsed_payloads:
+        if payload.get("type") != ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
+            continue
+        image_gen_usage = get_image_generation_usage(payload)
+        if image_gen_usage is not None:
+            return transform_image_usage(image_gen_usage)
+
+    for payload in reversed(parsed_payloads):
+        image_gen_usage = get_image_generation_usage(payload)
+        if image_gen_usage is not None and not is_zero_image_usage(image_gen_usage):
+            return transform_image_usage(image_gen_usage)
+    return None
+
+
+def get_parsed_payloads(raw_response: httpx.Response) -> List[dict]:
+    content_type = raw_response.headers.get("content-type", "")
+    body_text = raw_response.text or ""
+
+    if "text/event-stream" in content_type.lower() or looks_like_sse(body_text):
+        return parse_sse_payloads(body_text)
+
+    try:
+        response_json = raw_response.json()
+    except Exception:
+        response_json = {}
+    if isinstance(response_json, dict):
+        return [response_json]
+    return []
+
+
+def transform_image_usage(usage: dict) -> ImageUsage:
+    input_tokens_details = usage.get("input_tokens_details") or {}
+    return ImageUsage(
+        input_tokens=usage.get("input_tokens", 0),
+        input_tokens_details=ImageUsageInputTokensDetails(
+            image_tokens=input_tokens_details.get("image_tokens", 0),
+            text_tokens=input_tokens_details.get("text_tokens", 0),
+        ),
+        output_tokens=usage.get("output_tokens", 0),
+        total_tokens=usage.get("total_tokens", 0),
+    )
+
+
+def get_image_generation_usage(response_payload: Any) -> Optional[dict]:
+    if not isinstance(response_payload, dict):
+        return None
+
+    response = response_payload.get("response")
+    if isinstance(response, dict):
+        image_gen_usage = get_image_generation_usage(response)
+        if image_gen_usage is not None:
+            return image_gen_usage
+
+    tool_usage = response_payload.get("tool_usage")
+    if not isinstance(tool_usage, dict):
+        return None
+
+    image_gen_usage = tool_usage.get("image_gen")
+    if not isinstance(image_gen_usage, dict):
+        return None
+
+    input_tokens = image_gen_usage.get("input_tokens")
+    output_tokens = image_gen_usage.get("output_tokens")
+    if input_tokens is None or output_tokens is None:
+        return None
+
+    normalized_usage = dict(image_gen_usage)
+    if normalized_usage.get("total_tokens") is None:
+        normalized_usage["total_tokens"] = input_tokens + output_tokens
+    return normalized_usage
+
+
+def is_zero_image_usage(usage: dict) -> bool:
+    return (
+        (usage.get("input_tokens") or 0) == 0
+        and (usage.get("output_tokens") or 0) == 0
+        and (usage.get("total_tokens") or 0) == 0
+    )
+
+
+def looks_like_sse(body_text: str) -> bool:
+    trimmed_body = body_text.lstrip()
+    return (
+        trimmed_body.startswith("event:")
+        or trimmed_body.startswith("data:")
+        or "\nevent:" in body_text
+        or "\ndata:" in body_text
+    )
+
+
+def parse_sse_payloads(body_text: str) -> List[dict]:
+    payloads: List[dict] = []
+    for line in body_text.splitlines():
+        stripped_line = CustomStreamWrapper._strip_sse_data_from_chunk(line)
+        if not stripped_line:
+            continue
+        stripped_line = stripped_line.strip()
+        if not stripped_line or stripped_line == STREAM_SSE_DONE_STRING:
+            continue
+        try:
+            parsed = json.loads(stripped_line)
+        except json.JSONDecodeError:
+            continue
+        if isinstance(parsed, dict):
+            payloads.append(parsed)
+    return payloads
+
+
+def extract_images_from_payload(payload: dict) -> Tuple[List[str], List[str]]:
+    event_type = payload.get("type")
+    if event_type in (
+        ResponsesAPIStreamEvents.RESPONSE_FAILED,
+        ResponsesAPIStreamEvents.ERROR,
+    ):
+        error_obj = payload.get("error") or (payload.get("response") or {}).get("error")
+        raise OpenAIError(message=str(error_obj or payload), status_code=400)
+
+    partial_images: List[str] = []
+    if event_type in (
+        ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE,
+        "response.image_generation_call.partial_image",
+    ):
+        partial_image_b64 = payload.get("partial_image_b64")
+        b64_json = payload.get("b64_json")
+        if isinstance(partial_image_b64, str):
+            partial_images.append(partial_image_b64)
+        if isinstance(b64_json, str):
+            partial_images.append(b64_json)
+        return [], partial_images
+
+    candidates: List[str] = []
+    if event_type == "image_generation.completed":
+        b64_json = payload.get("b64_json")
+        if isinstance(b64_json, str):
+            candidates.append(b64_json)
+
+    response_payload = payload.get("response")
+    if isinstance(response_payload, dict):
+        candidates.extend(extract_images_from_nested_value(response_payload))
+
+    candidates.extend(extract_images_from_nested_value(payload))
+    return dedupe(candidates), dedupe(partial_images)
+
+
+def extract_images_from_nested_value(value: Any) -> List[str]:
+    images: List[str] = []
+    values_to_visit = [value]
+    visited_container_ids = set()
+
+    while values_to_visit:
+        current_value = values_to_visit.pop()
+        if isinstance(current_value, dict):
+            container_id = id(current_value)
+            if container_id in visited_container_ids:
+                continue
+            visited_container_ids.add(container_id)
+
+            value_type = current_value.get("type")
+            if value_type in ("image_generation_call", "image_generation"):
+                images.extend(get_image_strings_from_dict(current_value))
+            elif isinstance(current_value.get("b64_json"), str):
+                images.append(current_value["b64_json"])
+
+            values_to_visit.extend(reversed(list(current_value.values())))
+        elif isinstance(current_value, list):
+            container_id = id(current_value)
+            if container_id in visited_container_ids:
+                continue
+            visited_container_ids.add(container_id)
+
+            values_to_visit.extend(reversed(current_value))
+    return dedupe(images)
+
+
+def get_image_strings_from_dict(value: dict) -> List[str]:
+    images: List[str] = []
+    for key in ("result", "b64_json", "image"):
+        candidate = value.get(key)
+        if isinstance(candidate, str):
+            images.append(candidate)
+        elif isinstance(candidate, list):
+            images.extend(item for item in candidate if isinstance(item, str))
+    return images
+
+
+def dedupe(values: List[str]) -> List[str]:
+    seen = set()
+    deduped: List[str] = []
+    for value in values:
+        if value in seen:
+            continue
+        seen.add(value)
+        deduped.append(value)
+    return deduped

litellm/llms/chatgpt/image_generation/transformation.py

@@ -1,668 +0,0 @@
-import json
-import base64
-from io import BufferedReader, BytesIO
-from os import PathLike
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
-
-import httpx
-from httpx._types import RequestFiles
-
-from litellm.constants import STREAM_SSE_DONE_STRING
-from litellm.exceptions import AuthenticationError
-from litellm.images.utils import ImageEditRequestUtils
-from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
-from litellm.llms.base_llm.image_generation.transformation import (
-    BaseImageGenerationConfig,
-)
-from litellm.types.images.main import ImageEditOptionalRequestParams
-from litellm.llms.openai.common_utils import OpenAIError
-from litellm.types.llms.openai import (
-    AllMessageValues,
-    FileTypes,
-    OpenAIImageGenerationOptionalParams,
-    ResponsesAPIStreamEvents,
-)
-from litellm.types.router import GenericLiteLLMParams
-from litellm.types.utils import (
-    ImageObject,
-    ImageResponse,
-    ImageUsage,
-    ImageUsageInputTokensDetails,
-)
-from litellm.utils import CustomStreamWrapper
-
-from ..authenticator import Authenticator
-from ..common_utils import (
-    CHATGPT_API_BASE,
-    GetAccessTokenError,
-    ensure_chatgpt_session_id,
-    get_chatgpt_default_headers,
-    get_chatgpt_default_instructions,
-)
-
-GPT_IMAGE_MODEL_PREFIX = "gpt-image-"
-
-ALLOWED_OUTPUT_FORMATS = {"png", "jpeg", "webp"}
-INTERNAL_OPTIONAL_PARAMS = {"chatgpt_responses_model"}
-
-if TYPE_CHECKING:
-    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-
-
-class ChatGPTImageGenerationConfig(BaseImageGenerationConfig):
-    """
-    Bridge OpenAI-style Images API calls to ChatGPT/Codex Responses image generation.
-    """
-
-    def __init__(self) -> None:
-        self.authenticator = Authenticator()
-
-    def get_supported_openai_params(
-        self, model: str
-    ) -> List[OpenAIImageGenerationOptionalParams]:
-        return [
-            "output_format",
-            "size",
-        ]
-
-    def map_openai_params(
-        self,
-        non_default_params: dict,
-        optional_params: dict,
-        model: str,
-        drop_params: bool,
-    ) -> dict:
-        supported_params = self.get_supported_openai_params(model)
-        for key, value in non_default_params.items():
-            if key in optional_params:
-                continue
-            if key in supported_params:
-                optional_params[key] = value
-            elif drop_params:
-                continue
-            else:
-                raise ValueError(
-                    f"Parameter {key} is not supported for model {model}. "
-                    f"Supported parameters are {supported_params}. "
-                    "Set drop_params=True to drop unsupported parameters."
-                )
-        return optional_params
-
-    def validate_environment(
-        self,
-        headers: dict,
-        model: str,
-        messages: List[AllMessageValues],
-        optional_params: dict,
-        litellm_params: dict,
-        api_key: Optional[str] = None,
-        api_base: Optional[str] = None,
-    ) -> dict:
-        try:
-            access_token = self.authenticator.get_access_token()
-        except GetAccessTokenError as e:
-            raise AuthenticationError(
-                model=model,
-                llm_provider="chatgpt",
-                message=str(e),
-            )
-
-        account_id = self.authenticator.get_account_id()
-        session_id = ensure_chatgpt_session_id(litellm_params)
-        default_headers = get_chatgpt_default_headers(
-            access_token, account_id, session_id
-        )
-        return {**default_headers, **headers}
-
-    def get_complete_url(
-        self,
-        api_base: Optional[str],
-        api_key: Optional[str],
-        model: str,
-        optional_params: dict,
-        litellm_params: dict,
-        stream: Optional[bool] = None,
-    ) -> str:
-        api_base = self.authenticator.get_api_base() or CHATGPT_API_BASE
-        api_base = self._canonicalize_codex_api_base(api_base)
-        return f"{api_base}/responses"
-
-    @staticmethod
-    def _canonicalize_codex_api_base(api_base: str) -> str:
-        api_base = api_base.rstrip("/")
-        if api_base.endswith("/responses"):
-            api_base = api_base[: -len("/responses")]
-        if api_base.endswith("/backend-api"):
-            return f"{api_base}/codex"
-        return api_base
-
-    def transform_image_generation_request(
-        self,
-        model: str,
-        prompt: str,
-        optional_params: dict,
-        litellm_params: dict,
-        headers: dict,
-    ) -> dict:
-        self._validate_openai_image_generation_params(model, optional_params)
-
-        return self._build_responses_image_request(
-            model=model,
-            prompt=prompt,
-            optional_params=optional_params,
-            litellm_params=litellm_params,
-        )
-
-    def _build_responses_image_request(
-        self,
-        model: str,
-        prompt: Optional[str],
-        optional_params: dict,
-        litellm_params: dict,
-        input_images: Optional[List[Dict[str, Any]]] = None,
-    ) -> dict:
-        # Intentionally pinned fallback for ChatGPT image generation through the
-        # Codex Responses API. Users can override this per request or via
-        # litellm_params.
-        responses_model = (
-            optional_params.pop("chatgpt_responses_model", None)
-            or litellm_params.get("chatgpt_responses_model")
-            or "gpt-5.5"
-        )
-        content: List[Dict[str, Any]] = []
-        if prompt:
-            content.append({"type": "input_text", "text": prompt})
-        if input_images:
-            content.extend(input_images)
-
-        request: Dict[str, Any] = {
-            "model": responses_model,
-            "input": [
-                {
-                    "role": "user",
-                    "content": content,
-                }
-            ],
-            "instructions": get_chatgpt_default_instructions(),
-            "tools": [{"type": "image_generation", "model": model}],
-            "tool_choice": {"type": "image_generation"},
-            "stream": True,
-            "store": False,
-        }
-
-        image_tool = request["tools"][0]
-        for key in (
-            "output_format",
-            "size",
-        ):
-            if optional_params.get(key) is not None:
-                image_tool[key] = optional_params[key]
-
-        return request
-
-    def _validate_openai_image_generation_params(
-        self, model: str, optional_params: dict
-    ) -> None:
-        if not model.startswith(GPT_IMAGE_MODEL_PREFIX):
-            raise ValueError(
-                "ChatGPT image generation requires a GPT Image model "
-                "(for example gpt-image-1.5 or gpt-image-2)."
-            )
-
-        supported_params = set(self.get_supported_openai_params(model))
-        unsupported_params = [
-            key
-            for key in optional_params
-            if key not in supported_params and key not in INTERNAL_OPTIONAL_PARAMS
-        ]
-        if unsupported_params:
-            raise ValueError(
-                f"Parameters {unsupported_params} are not supported for model {model}. "
-                f"Supported parameters are {sorted(supported_params)}."
-            )
-
-        output_format = optional_params.get("output_format")
-        if output_format is not None and output_format not in ALLOWED_OUTPUT_FORMATS:
-            raise ValueError("output_format must be one of png, jpeg, or webp")
-
-    def transform_image_generation_response(
-        self,
-        model: str,
-        raw_response: httpx.Response,
-        model_response: ImageResponse,
-        logging_obj: "LiteLLMLoggingObj",
-        request_data: dict,
-        optional_params: dict,
-        litellm_params: dict,
-        encoding: Any,
-        api_key: Optional[str] = None,
-        json_mode: Optional[bool] = None,
-    ) -> ImageResponse:
-        logging_obj.post_call(
-            input=request_data.get("input", ""),
-            api_key=api_key,
-            additional_args={"complete_input_dict": request_data},
-            original_response=raw_response.text,
-        )
-
-        image_payloads = self._extract_image_payloads(raw_response)
-        if not image_payloads:
-            raise OpenAIError(
-                message="No image data found in ChatGPT image generation response",
-                status_code=raw_response.status_code,
-            )
-
-        response = ImageResponse(
-            data=[
-                ImageObject(b64_json=image_payload) for image_payload in image_payloads
-            ]
-        )
-        response.usage = None
-        image_usage = self._extract_image_usage(raw_response)
-        if image_usage is not None:
-            response.usage = image_usage
-        response.size = optional_params.get("size")
-        response.output_format = optional_params.get("output_format")
-        response._hidden_params["model"] = model
-        return response
-
-    def _extract_image_payloads(self, raw_response: httpx.Response) -> List[str]:
-        content_type = raw_response.headers.get("content-type", "")
-        body_text = raw_response.text or ""
-        parsed_payloads: List[dict] = []
-
-        if "text/event-stream" in content_type.lower() or self._looks_like_sse(
-            body_text
-        ):
-            parsed_payloads = self._parse_sse_payloads(body_text)
-        else:
-            try:
-                response_json = raw_response.json()
-            except Exception:
-                response_json = {}
-            if isinstance(response_json, dict):
-                parsed_payloads = [response_json]
-
-        images: List[str] = []
-        partial_images: List[str] = []
-        for payload in parsed_payloads:
-            extracted_images, extracted_partial_images = (
-                self._extract_images_from_payload(payload)
-            )
-            images.extend(extracted_images)
-            partial_images.extend(extracted_partial_images)
-        return self._dedupe(images) or self._dedupe(partial_images)
-
-    def _extract_image_usage(
-        self, raw_response: httpx.Response
-    ) -> Optional[ImageUsage]:
-        parsed_payloads = self._get_parsed_payloads(raw_response)
-
-        for payload in parsed_payloads:
-            if payload.get("type") != ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
-                continue
-            image_gen_usage = self._get_image_generation_usage(payload)
-            if image_gen_usage is not None:
-                return self._transform_image_usage(image_gen_usage)
-
-        for payload in reversed(parsed_payloads):
-            image_gen_usage = self._get_image_generation_usage(payload)
-            if image_gen_usage is not None and not self._is_zero_image_usage(
-                image_gen_usage
-            ):
-                return self._transform_image_usage(image_gen_usage)
-        return None
-
-    def _get_parsed_payloads(self, raw_response: httpx.Response) -> List[dict]:
-        content_type = raw_response.headers.get("content-type", "")
-        body_text = raw_response.text or ""
-
-        if "text/event-stream" in content_type.lower() or self._looks_like_sse(
-            body_text
-        ):
-            return self._parse_sse_payloads(body_text)
-
-        try:
-            response_json = raw_response.json()
-        except Exception:
-            response_json = {}
-        if isinstance(response_json, dict):
-            return [response_json]
-        return []
-
-    @staticmethod
-    def _transform_image_usage(usage: dict) -> ImageUsage:
-        input_tokens_details = usage.get("input_tokens_details") or {}
-        return ImageUsage(
-            input_tokens=usage.get("input_tokens", 0),
-            input_tokens_details=ImageUsageInputTokensDetails(
-                image_tokens=input_tokens_details.get("image_tokens", 0),
-                text_tokens=input_tokens_details.get("text_tokens", 0),
-            ),
-            output_tokens=usage.get("output_tokens", 0),
-            total_tokens=usage.get("total_tokens", 0),
-        )
-
-    @staticmethod
-    def _get_image_generation_usage(response_payload: Any) -> Optional[dict]:
-        if not isinstance(response_payload, dict):
-            return None
-
-        response = response_payload.get("response")
-        if isinstance(response, dict):
-            image_gen_usage = ChatGPTImageGenerationConfig._get_image_generation_usage(
-                response
-            )
-            if image_gen_usage is not None:
-                return image_gen_usage
-
-        tool_usage = response_payload.get("tool_usage")
-        if not isinstance(tool_usage, dict):
-            return None
-
-        image_gen_usage = tool_usage.get("image_gen")
-        if not isinstance(image_gen_usage, dict):
-            return None
-
-        input_tokens = image_gen_usage.get("input_tokens")
-        output_tokens = image_gen_usage.get("output_tokens")
-        if input_tokens is None or output_tokens is None:
-            return None
-
-        normalized_usage = dict(image_gen_usage)
-        if normalized_usage.get("total_tokens") is None:
-            normalized_usage["total_tokens"] = input_tokens + output_tokens
-        return normalized_usage
-
-    @staticmethod
-    def _is_zero_image_usage(usage: dict) -> bool:
-        return (
-            (usage.get("input_tokens") or 0) == 0
-            and (usage.get("output_tokens") or 0) == 0
-            and (usage.get("total_tokens") or 0) == 0
-        )
-
-    @staticmethod
-    def _looks_like_sse(body_text: str) -> bool:
-        trimmed_body = body_text.lstrip()
-        return (
-            trimmed_body.startswith("event:")
-            or trimmed_body.startswith("data:")
-            or "\nevent:" in body_text
-            or "\ndata:" in body_text
-        )
-
-    @staticmethod
-    def _parse_sse_payloads(body_text: str) -> List[dict]:
-        payloads: List[dict] = []
-        for line in body_text.splitlines():
-            stripped_line = CustomStreamWrapper._strip_sse_data_from_chunk(line)
-            if not stripped_line:
-                continue
-            stripped_line = stripped_line.strip()
-            if not stripped_line or stripped_line == STREAM_SSE_DONE_STRING:
-                continue
-            try:
-                parsed = json.loads(stripped_line)
-            except json.JSONDecodeError:
-                continue
-            if isinstance(parsed, dict):
-                payloads.append(parsed)
-        return payloads
-
-    def _extract_images_from_payload(
-        self, payload: dict
-    ) -> Tuple[List[str], List[str]]:
-        event_type = payload.get("type")
-        if event_type in (
-            ResponsesAPIStreamEvents.RESPONSE_FAILED,
-            ResponsesAPIStreamEvents.ERROR,
-        ):
-            error_obj = payload.get("error") or (payload.get("response") or {}).get(
-                "error"
-            )
-            raise OpenAIError(message=str(error_obj or payload), status_code=400)
-
-        partial_images: List[str] = []
-        if event_type in (
-            ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE,
-            "response.image_generation_call.partial_image",
-        ):
-            partial_image_b64 = payload.get("partial_image_b64")
-            b64_json = payload.get("b64_json")
-            if isinstance(partial_image_b64, str):
-                partial_images.append(partial_image_b64)
-            if isinstance(b64_json, str):
-                partial_images.append(b64_json)
-            return [], partial_images
-
-        candidates: List[str] = []
-        if event_type == "image_generation.completed":
-            b64_json = payload.get("b64_json")
-            if isinstance(b64_json, str):
-                candidates.append(b64_json)
-
-        response_payload = payload.get("response")
-        if isinstance(response_payload, dict):
-            candidates.extend(self._extract_images_from_nested_value(response_payload))
-
-        candidates.extend(self._extract_images_from_nested_value(payload))
-        return self._dedupe(candidates), self._dedupe(partial_images)
-
-    def _extract_images_from_nested_value(self, value: Any) -> List[str]:
-        images: List[str] = []
-        values_to_visit = [value]
-        visited_container_ids = set()
-
-        while values_to_visit:
-            current_value = values_to_visit.pop()
-            if isinstance(current_value, dict):
-                container_id = id(current_value)
-                if container_id in visited_container_ids:
-                    continue
-                visited_container_ids.add(container_id)
-
-                value_type = current_value.get("type")
-                if value_type in ("image_generation_call", "image_generation"):
-                    images.extend(self._get_image_strings_from_dict(current_value))
-                elif isinstance(current_value.get("b64_json"), str):
-                    images.append(current_value["b64_json"])
-
-                values_to_visit.extend(reversed(list(current_value.values())))
-            elif isinstance(current_value, list):
-                container_id = id(current_value)
-                if container_id in visited_container_ids:
-                    continue
-                visited_container_ids.add(container_id)
-
-                values_to_visit.extend(reversed(current_value))
-        return self._dedupe(images)
-
-    @staticmethod
-    def _get_image_strings_from_dict(value: dict) -> List[str]:
-        images: List[str] = []
-        for key in ("result", "b64_json", "image"):
-            candidate = value.get(key)
-            if isinstance(candidate, str):
-                images.append(candidate)
-            elif isinstance(candidate, list):
-                images.extend(item for item in candidate if isinstance(item, str))
-        return images
-
-    @staticmethod
-    def _dedupe(values: List[str]) -> List[str]:
-        seen = set()
-        deduped: List[str] = []
-        for value in values:
-            if value in seen:
-                continue
-            seen.add(value)
-            deduped.append(value)
-        return deduped
-
-    def get_error_class(
-        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
-    ) -> OpenAIError:
-        return OpenAIError(
-            message=error_message,
-            status_code=status_code,
-            headers=headers,
-        )
-
-
-class ChatGPTImageEditConfig(BaseImageEditConfig):
-    """
-    Bridge OpenAI-style Images Edits calls to ChatGPT/Codex Responses image generation.
-    """
-
-    def __init__(self) -> None:
-        self.image_generation_config = ChatGPTImageGenerationConfig()
-
-    def get_supported_openai_params(self, model: str) -> List[str]:
-        return ["size"]
-
-    def map_openai_params(
-        self,
-        image_edit_optional_params: ImageEditOptionalRequestParams,
-        model: str,
-        drop_params: bool,
-    ) -> Dict[str, Any]:
-        supported_params = self.get_supported_openai_params(model)
-        return {
-            key: value
-            for key, value in image_edit_optional_params.items()
-            if key in supported_params
-        }
-
-    def validate_environment(
-        self,
-        headers: dict,
-        model: str,
-        api_key: Optional[str] = None,
-        litellm_params: Optional[dict] = None,
-        api_base: Optional[str] = None,
-    ) -> dict:
-        return self.image_generation_config.validate_environment(
-            headers=headers,
-            model=model,
-            messages=[],
-            optional_params={},
-            litellm_params=litellm_params or {},
-            api_key=api_key,
-            api_base=api_base,
-        )
-
-    def get_complete_url(
-        self,
-        model: str,
-        api_base: Optional[str],
-        litellm_params: dict,
-    ) -> str:
-        return self.image_generation_config.get_complete_url(
-            api_base=api_base,
-            api_key=litellm_params.get("api_key"),
-            model=model,
-            optional_params={},
-            litellm_params=litellm_params,
-        )
-
-    def use_multipart_form_data(self) -> bool:
-        return False
-
-    def transform_image_edit_request(
-        self,
-        model: str,
-        prompt: Optional[str],
-        image: Optional[FileTypes],
-        image_edit_optional_request_params: Dict,
-        litellm_params: GenericLiteLLMParams,
-        headers: dict,
-    ) -> Tuple[Dict[str, Any], RequestFiles]:
-        optional_params = dict(image_edit_optional_request_params)
-        self.image_generation_config._validate_openai_image_generation_params(
-            model, optional_params
-        )
-
-        input_images = self._prepare_input_images(image)
-        if not input_images:
-            raise ValueError("ChatGPT image edit requires at least one image.")
-
-        request = self.image_generation_config._build_responses_image_request(
-            model=model,
-            prompt=prompt,
-            optional_params=optional_params,
-            litellm_params=dict(litellm_params),
-            input_images=input_images,
-        )
-        return request, []
-
-    def transform_image_edit_response(
-        self,
-        model: str,
-        raw_response: httpx.Response,
-        logging_obj: "LiteLLMLoggingObj",
-    ) -> ImageResponse:
-        return self.image_generation_config.transform_image_generation_response(
-            model=model,
-            raw_response=raw_response,
-            model_response=ImageResponse(),
-            logging_obj=logging_obj,
-            request_data={},
-            optional_params={},
-            litellm_params={},
-            encoding=None,
-        )
-
-    def _prepare_input_images(
-        self, image: Optional[Union[FileTypes, List[FileTypes]]]
-    ) -> List[Dict[str, Any]]:
-        if image is None:
-            return []
-
-        images = image if isinstance(image, list) else [image]
-        input_images: List[Dict[str, Any]] = []
-        for img in images:
-            if img is None:
-                continue
-            mime_type = ImageEditRequestUtils.get_image_content_type(img)
-            image_bytes = self._read_image_bytes(img)
-            b64_data = base64.b64encode(image_bytes).decode("utf-8")
-            input_images.append(
-                {
-                    "type": "input_image",
-                    "image_url": f"data:{mime_type};base64,{b64_data}",
-                }
-            )
-        return input_images
-
-    @staticmethod
-    def _read_image_bytes(image: FileTypes) -> bytes:
-        if isinstance(image, bytes):
-            return image
-        if isinstance(image, BytesIO):
-            current_pos = image.tell()
-            image.seek(0)
-            data = image.read()
-            image.seek(current_pos)
-            return data
-        if isinstance(image, BufferedReader):
-            current_pos = image.tell()
-            image.seek(0)
-            data = image.read()
-            image.seek(current_pos)
-            return data
-        if isinstance(image, tuple):
-            return ChatGPTImageEditConfig._read_image_bytes(image[1])
-        if isinstance(image, PathLike):
-            with open(image, "rb") as image_file:
-                return image_file.read()
-        raise ValueError("Unsupported image type for ChatGPT image edit.")
-
-    def get_error_class(
-        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
-    ) -> OpenAIError:
-        return self.image_generation_config.get_error_class(
-            error_message=error_message,
-            status_code=status_code,
-            headers=headers,
-        )

litellm/utils.py

@@ -9123,7 +9123,7 @@ class ProviderConfigManager:
 
             return get_openai_image_edit_config(model=model)
         elif LlmProviders.CHATGPT == provider:
-            from litellm.llms.chatgpt.image_generation import ChatGPTImageEditConfig
+            from litellm.llms.chatgpt.image_edit import ChatGPTImageEditConfig
 
             return ChatGPTImageEditConfig()
         elif LlmProviders.AZURE == provider:

tests/test_litellm/llms/chatgpt/chatgpt_image_test_utils.py

@@ -0,0 +1,10 @@
+from typing import Any
+
+
+class MockLogging:
+    def post_call(self, *args, **kwargs):
+        pass
+
+
+def mock_logging() -> Any:
+    return MockLogging()

tests/test_litellm/llms/chatgpt/image_edit/test_transformation.py

@@ -0,0 +1,189 @@
+from io import BytesIO
+from typing import Any, cast
+
+import httpx
+import pytest
+
+from litellm.llms.chatgpt.image_edit import ChatGPTImageEditConfig
+from litellm.llms.openai.common_utils import OpenAIError
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+from tests.test_litellm.llms.chatgpt.chatgpt_image_test_utils import mock_logging
+
+
+@pytest.fixture(autouse=True)
+def _chatgpt_token_dir(monkeypatch, tmp_path):
+    monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
+
+
+def test_chatgpt_image_edit_config_registered():
+    config = ProviderConfigManager.get_provider_image_edit_config(
+        model="gpt-image-2",
+        provider=LlmProviders.CHATGPT,
+    )
+
+    assert isinstance(config, ChatGPTImageEditConfig)
+
+
+def test_chatgpt_image_edit_transforms_request():
+    config = ChatGPTImageEditConfig()
+    png_bytes = (
+        b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
+        b"\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89"
+    )
+
+    request, files = config.transform_image_edit_request(
+        model="gpt-image-2",
+        prompt="replace the background with a warm sunset",
+        image=png_bytes,
+        image_edit_optional_request_params={"size": "1024x1024"},
+        litellm_params=cast(Any, {"chatgpt_responses_model": "gpt-5.5"}),
+        headers={},
+    )
+
+    assert files == []
+    assert request["model"] == "gpt-5.5"
+    assert request["input"][0]["content"][0] == {
+        "type": "input_text",
+        "text": "replace the background with a warm sunset",
+    }
+    assert request["input"][0]["content"][1]["type"] == "input_image"
+    assert request["input"][0]["content"][1]["image_url"].startswith(
+        "data:image/png;base64,"
+    )
+    assert request["tools"] == [
+        {
+            "type": "image_generation",
+            "model": "gpt-image-2",
+            "size": "1024x1024",
+        }
+    ]
+    assert request["tool_choice"] == {"type": "image_generation"}
+
+
+def test_chatgpt_image_edit_uses_json_requests():
+    config = ChatGPTImageEditConfig()
+
+    assert config.use_multipart_form_data() is False
+
+
+def test_chatgpt_image_edit_requires_image():
+    config = ChatGPTImageEditConfig()
+
+    with pytest.raises(ValueError, match="requires at least one image"):
+        config.transform_image_edit_request(
+            model="gpt-image-2",
+            prompt="edit this image",
+            image=None,
+            image_edit_optional_request_params={},
+            litellm_params=cast(Any, {}),
+            headers={},
+        )
+
+
+def test_chatgpt_image_edit_delegates_environment_and_url():
+    config = ChatGPTImageEditConfig()
+
+    class FakeImageGenerationConfig:
+        def validate_environment(self, **kwargs):
+            assert kwargs["messages"] == []
+            assert kwargs["optional_params"] == {}
+            assert kwargs["litellm_params"] == {"session_id": "session-123"}
+            assert kwargs["api_key"] == "api-key"
+            assert kwargs["api_base"] == "https://ignored.test"
+            return {"Authorization": "Bearer token"}
+
+        def get_complete_url(self, **kwargs):
+            assert kwargs["api_key"] == "api-key"
+            assert kwargs["optional_params"] == {}
+            return "https://chatgpt.com/backend-api/codex/responses"
+
+    config.image_generation_config = cast(Any, FakeImageGenerationConfig())
+
+    assert config.get_supported_openai_params("gpt-image-2") == ["size"]
+    assert config.map_openai_params(
+        image_edit_optional_params={"size": "1024x1024", "quality": "high"},
+        model="gpt-image-2",
+        drop_params=False,
+    ) == {"size": "1024x1024"}
+    assert config.validate_environment(
+        headers={},
+        model="gpt-image-2",
+        api_key="api-key",
+        litellm_params={"session_id": "session-123"},
+        api_base="https://ignored.test",
+    ) == {"Authorization": "Bearer token"}
+    assert (
+        config.get_complete_url(
+            model="gpt-image-2",
+            api_base="https://ignored.test",
+            litellm_params={"api_key": "api-key"},
+        )
+        == "https://chatgpt.com/backend-api/codex/responses"
+    )
+
+
+def test_chatgpt_image_edit_transform_response_and_error_class():
+    config = ChatGPTImageEditConfig()
+
+    raw_response = httpx.Response(
+        status_code=200,
+        json={
+            "output": [
+                {
+                    "type": "image_generation",
+                    "image": "edited-image-data",
+                }
+            ]
+        },
+    )
+
+    response = config.transform_image_edit_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        logging_obj=mock_logging(),
+    )
+    error = config.get_error_class(
+        error_message="bad edit",
+        status_code=400,
+        headers={"x-request-id": "req-456"},
+    )
+
+    assert response.data is not None
+    assert response.data[0].b64_json == "edited-image-data"
+    assert isinstance(error, OpenAIError)
+    assert error.status_code == 400
+    assert error.message == "bad edit"
+
+
+def test_chatgpt_image_edit_prepare_input_images_handles_supported_file_types(
+    tmp_path,
+):
+    config = ChatGPTImageEditConfig()
+    image_path = tmp_path / "image.png"
+    image_path.write_bytes(b"path-bytes")
+
+    bytes_io = BytesIO(b"bytes-io-data")
+    bytes_io.seek(5)
+    with image_path.open("rb") as buffered_reader:
+        buffered_reader.seek(2)
+        input_images = config._prepare_input_images(
+            [
+                None,
+                bytes_io,
+                buffered_reader,
+                ("image.png", b"tuple-bytes", "image/png"),
+                image_path,
+            ]
+        )
+        assert buffered_reader.tell() == 2
+
+    assert bytes_io.tell() == 5
+    assert [image["type"] for image in input_images] == ["input_image"] * 4
+    assert input_images[0]["image_url"].startswith("data:image/png;base64,")
+    assert input_images[1]["image_url"].startswith("data:image/png;base64,")
+    assert input_images[2]["image_url"].startswith("data:image/png;base64,")
+    assert input_images[3]["image_url"].startswith("data:image/png;base64,")
+
+    with pytest.raises(ValueError, match="Unsupported image type"):
+        config._read_image_bytes(cast(Any, object()))

tests/test_litellm/llms/chatgpt/image_generation/test_request_transformation.py

@@ -0,0 +1,395 @@
+from typing import Any, cast
+
+import pytest
+
+from litellm.exceptions import AuthenticationError
+from litellm.llms.chatgpt.common_utils import GetAccessTokenError
+from litellm.llms.chatgpt.image_generation import ChatGPTImageGenerationConfig
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+
+
+@pytest.fixture(autouse=True)
+def _chatgpt_token_dir(monkeypatch, tmp_path):
+    monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
+
+
+def test_chatgpt_image_generation_transforms_request():
+    config = ChatGPTImageGenerationConfig()
+
+    request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a quiet harbor at sunrise",
+        optional_params={"size": "1024x1024", "output_format": "png"},
+        litellm_params={"chatgpt_responses_model": "gpt-5.5"},
+        headers={},
+    )
+
+    assert request["model"] == "gpt-5.5"
+    assert request["input"] == [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "input_text",
+                    "text": "draw a quiet harbor at sunrise",
+                }
+            ],
+        }
+    ]
+    assert request["stream"] is True
+    assert request["store"] is False
+    assert request["tools"] == [
+        {
+            "type": "image_generation",
+            "model": "gpt-image-2",
+            "size": "1024x1024",
+            "output_format": "png",
+        }
+    ]
+    assert request["tool_choice"] == {"type": "image_generation"}
+
+
+def test_chatgpt_image_generation_does_not_add_openai_defaults():
+    config = ChatGPTImageGenerationConfig()
+
+    request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a quiet harbor at sunrise",
+        optional_params={},
+        litellm_params={"chatgpt_responses_model": "gpt-5.5"},
+        headers={},
+    )
+
+    assert request["tools"] == [{"type": "image_generation", "model": "gpt-image-2"}]
+
+
+def test_chatgpt_image_generation_forwards_supported_generate_params():
+    config = ChatGPTImageGenerationConfig()
+
+    request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a quiet harbor at sunrise",
+        optional_params={
+            "output_format": "webp",
+            "size": "1536x1024",
+        },
+        litellm_params={"chatgpt_responses_model": "gpt-5.5"},
+        headers={},
+    )
+
+    assert request["tools"] == [
+        {
+            "type": "image_generation",
+            "model": "gpt-image-2",
+            "output_format": "webp",
+            "size": "1536x1024",
+        }
+    ]
+
+
+@pytest.mark.parametrize(
+    "optional_params, error",
+    [
+        ({"output_format": "jpg"}, "output_format must be one of png, jpeg, or webp"),
+    ],
+)
+def test_chatgpt_image_generation_validates_params(optional_params, error):
+    config = ChatGPTImageGenerationConfig()
+
+    with pytest.raises(ValueError, match=error):
+        config.transform_image_generation_request(
+            model="gpt-image-2",
+            prompt="draw a quiet harbor at sunrise",
+            optional_params=optional_params,
+            litellm_params={"chatgpt_responses_model": "gpt-5.5"},
+            headers={},
+        )
+
+
+def test_chatgpt_image_generation_config_registered():
+    config = ProviderConfigManager.get_provider_image_generation_config(
+        model="gpt-image-2",
+        provider=LlmProviders.CHATGPT,
+    )
+
+    assert isinstance(config, ChatGPTImageGenerationConfig)
+
+
+def test_chatgpt_image_generation_only_supports_prompt_output_format_and_size():
+    config = ChatGPTImageGenerationConfig()
+
+    assert config.get_supported_openai_params("gpt-image-2") == [
+        "output_format",
+        "size",
+    ]
+
+
+def test_chatgpt_image_generation_rejects_unsupported_optional_params():
+    config = ChatGPTImageGenerationConfig()
+
+    with pytest.raises(
+        ValueError, match="Parameters \\['quality'\\] are not supported"
+    ):
+        config.transform_image_generation_request(
+            model="gpt-image-2",
+            prompt="draw a cat",
+            optional_params={"quality": "high"},
+            litellm_params={},
+            headers={},
+        )
+
+
+def test_chatgpt_image_generation_maps_supported_openai_params():
+    config = ChatGPTImageGenerationConfig()
+
+    optional_params = {"output_format": "png"}
+    result = config.map_openai_params(
+        non_default_params={
+            "quality": "low",
+            "size": "1024x1024",
+            "unsupported": "drop-me",
+        },
+        optional_params=optional_params,
+        model="gpt-image-2",
+        drop_params=True,
+    )
+
+    assert result is optional_params
+    assert result == {"output_format": "png", "size": "1024x1024"}
+
+
+def test_chatgpt_image_generation_rejects_unsupported_openai_param():
+    config = ChatGPTImageGenerationConfig()
+
+    with pytest.raises(ValueError, match="Parameter unsupported is not supported"):
+        config.map_openai_params(
+            non_default_params={"unsupported": "keep-me"},
+            optional_params={},
+            model="gpt-image-2",
+            drop_params=False,
+        )
+
+
+def test_chatgpt_image_generation_validates_environment():
+    config = ChatGPTImageGenerationConfig()
+
+    class FakeAuthenticator:
+        def get_access_token(self):
+            return "access-token"
+
+        def get_account_id(self):
+            return "account-id"
+
+    config.authenticator = cast(Any, FakeAuthenticator())
+
+    headers = config.validate_environment(
+        headers={"content-type": "application/custom", "x-extra": "1"},
+        model="gpt-image-2",
+        messages=[],
+        optional_params={},
+        litellm_params={"session_id": "session-123"},
+    )
+
+    assert headers["Authorization"] == "Bearer access-token"
+    assert headers["ChatGPT-Account-Id"] == "account-id"
+    assert headers["session_id"] == "session-123"
+    assert headers["content-type"] == "application/custom"
+    assert headers["x-extra"] == "1"
+
+
+def test_chatgpt_image_generation_validate_environment_auth_error():
+    config = ChatGPTImageGenerationConfig()
+
+    class FakeAuthenticator:
+        def get_access_token(self):
+            raise GetAccessTokenError(status_code=401, message="token expired")
+
+    config.authenticator = cast(Any, FakeAuthenticator())
+
+    with pytest.raises(AuthenticationError, match="token expired"):
+        config.validate_environment(
+            headers={},
+            model="gpt-image-2",
+            messages=[],
+            optional_params={},
+            litellm_params={},
+        )
+
+
+@pytest.mark.parametrize(
+    "server_api_base, expected",
+    [
+        (
+            "https://chatgpt.com/backend-api",
+            "https://chatgpt.com/backend-api/codex/responses",
+        ),
+        (
+            "https://chatgpt.com/backend-api/responses",
+            "https://chatgpt.com/backend-api/codex/responses",
+        ),
+        (
+            "https://example.test/custom/",
+            "https://example.test/custom/responses",
+        ),
+    ],
+)
+def test_chatgpt_image_generation_get_complete_url_canonicalizes_server_api_base(
+    server_api_base, expected
+):
+    config = ChatGPTImageGenerationConfig()
+
+    class FakeAuthenticator:
+        def get_api_base(self):
+            return server_api_base
+
+    config.authenticator = cast(Any, FakeAuthenticator())
+
+    assert (
+        config.get_complete_url(
+            api_base=None,
+            api_key=None,
+            model="gpt-image-2",
+            optional_params={},
+            litellm_params={},
+        )
+        == expected
+    )
+
+
+def test_chatgpt_image_generation_get_complete_url_ignores_request_api_base():
+    config = ChatGPTImageGenerationConfig()
+
+    class FakeAuthenticator:
+        def get_api_base(self):
+            return "https://chatgpt.com/backend-api"
+
+    config.authenticator = cast(Any, FakeAuthenticator())
+
+    assert (
+        config.get_complete_url(
+            api_base="https://attacker.test/collect",
+            api_key=None,
+            model="gpt-image-2",
+            optional_params={},
+            litellm_params={},
+        )
+        == "https://chatgpt.com/backend-api/codex/responses"
+    )
+
+
+def test_chatgpt_image_generation_get_complete_url_uses_authenticator_api_base():
+    config = ChatGPTImageGenerationConfig()
+
+    class FakeAuthenticator:
+        def get_api_base(self):
+            return "https://example.test/backend-api"
+
+    config.authenticator = cast(Any, FakeAuthenticator())
+
+    assert (
+        config.get_complete_url(
+            api_base=None,
+            api_key=None,
+            model="gpt-image-2",
+            optional_params={},
+            litellm_params={},
+        )
+        == "https://example.test/backend-api/codex/responses"
+    )
+
+
+def test_chatgpt_image_generation_uses_optional_responses_model():
+    config = ChatGPTImageGenerationConfig()
+    optional_params = {"chatgpt_responses_model": "gpt-override"}
+
+    request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a cat",
+        optional_params=optional_params,
+        litellm_params={"chatgpt_responses_model": "gpt-litellm-param"},
+        headers={},
+    )
+
+    assert request["model"] == "gpt-override"
+    assert "chatgpt_responses_model" not in optional_params
+
+    litellm_params_request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a cat",
+        optional_params={},
+        litellm_params={"chatgpt_responses_model": "gpt-litellm-param"},
+        headers={},
+    )
+    assert litellm_params_request["model"] == "gpt-litellm-param"
+
+    default_request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a cat",
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )
+    assert default_request["model"] == "gpt-5.5"
+
+
+@pytest.mark.parametrize(
+    "model, optional_params, error",
+    [
+        ("dall-e-3", {}, "requires a GPT Image model"),
+        ("gpt-image-1.5", {"size": "auto"}, None),
+        ("gpt-image-2", {"size": "auto"}, None),
+        ("gpt-image-2", {"size": "bad-size"}, None),
+    ],
+)
+def test_chatgpt_image_generation_validates_additional_param_paths(
+    model, optional_params, error
+):
+    config = ChatGPTImageGenerationConfig()
+
+    if error is None:
+        config.transform_image_generation_request(
+            model=model,
+            prompt="draw a cat",
+            optional_params=optional_params,
+            litellm_params={},
+            headers={},
+        )
+        return
+
+    with pytest.raises(ValueError, match=error):
+        config.transform_image_generation_request(
+            model=model,
+            prompt="draw a cat",
+            optional_params=optional_params,
+            litellm_params={},
+            headers={},
+        )
+
+
+def test_chatgpt_image_generation_forwards_size_without_local_constraints():
+    config = ChatGPTImageGenerationConfig()
+
+    request = config.transform_image_generation_request(
+        model="gpt-image-2",
+        prompt="draw a cat",
+        optional_params={"size": "bad-size"},
+        litellm_params={},
+        headers={},
+    )
+
+    assert request["tools"][0]["size"] == "bad-size"
+
+
+def test_chatgpt_image_generation_map_openai_params_keeps_existing_value():
+    config = ChatGPTImageGenerationConfig()
+
+    optional_params = {"size": "1024x1024"}
+    result = config.map_openai_params(
+        non_default_params={"size": "1536x1024"},
+        optional_params=optional_params,
+        model="gpt-image-2",
+        drop_params=False,
+    )
+
+    assert result == {"size": "1024x1024"}

tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py

@@ -0,0 +1,211 @@
+import httpx
+import pytest
+
+from litellm.llms.chatgpt.image_generation import ChatGPTImageGenerationConfig
+from litellm.llms.openai.common_utils import OpenAIError
+from litellm.types.utils import ImageResponse
+from tests.test_litellm.llms.chatgpt.chatgpt_image_test_utils import mock_logging
+
+
+@pytest.fixture(autouse=True)
+def _chatgpt_token_dir(monkeypatch, tmp_path):
+    monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
+
+
+def test_chatgpt_image_generation_extracts_b64_from_sse_completed_response():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "text/event-stream"},
+        text=(
+            'data: {"type":"response.completed","response":{"output":['
+            '{"type":"image_generation_call","result":"b64-image-data"}]}}\n\n'
+            "data: [DONE]\n\n"
+        ),
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={"input": "draw a cat"},
+        optional_params={"size": "1024x1024"},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.data is not None
+    assert response.data[0].b64_json == "b64-image-data"
+    assert response.size == "1024x1024"
+    assert response.quality is None
+    assert response.output_format is None
+    assert response.usage is None
+    assert response._hidden_params is not None
+    assert response._hidden_params["model"] == "gpt-image-2"
+
+
+def test_chatgpt_image_generation_extracts_b64_from_deep_nested_payload():
+    config = ChatGPTImageGenerationConfig()
+    nested_payload = {"type": "image_generation_call", "result": "b64-image-data"}
+    for _ in range(1200):
+        nested_payload = {"nested": [nested_payload]}
+
+    images, partial_images = config._extract_images_from_payload(nested_payload)
+
+    assert images == ["b64-image-data"]
+    assert partial_images == []
+
+
+def test_chatgpt_image_generation_extracts_json_response():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        json={
+            "output": [
+                {
+                    "type": "image_generation",
+                    "image": ["b64-image-data", "b64-image-data", 123],
+                }
+            ],
+            "tool_usage": {
+                "image_gen": {
+                    "input_tokens": 11,
+                    "input_tokens_details": {"image_tokens": 1, "text_tokens": 10},
+                    "output_tokens": 22,
+                }
+            },
+        },
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={},
+        optional_params={"output_format": "png"},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.data is not None
+    assert [item.b64_json for item in response.data] == ["b64-image-data"]
+    assert response.output_format == "png"
+    assert response.usage is not None
+    assert response.usage.input_tokens == 11
+    assert response.usage.input_tokens_details.image_tokens == 1
+    assert response.usage.input_tokens_details.text_tokens == 10
+    assert response.usage.output_tokens == 22
+    assert response.usage.total_tokens == 33
+
+
+def test_chatgpt_image_generation_raises_when_no_image_data():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(status_code=200, json={"output": []})
+
+    with pytest.raises(OpenAIError, match="No image data found"):
+        config.transform_image_generation_response(
+            model="gpt-image-2",
+            raw_response=raw_response,
+            model_response=ImageResponse(),
+            logging_obj=mock_logging(),
+            request_data={},
+            optional_params={},
+            litellm_params={},
+            encoding=None,
+        )
+
+
+def test_chatgpt_image_generation_raises_provider_error_event():
+    config = ChatGPTImageGenerationConfig()
+
+    with pytest.raises(OpenAIError, match="image blocked"):
+        config._extract_images_from_payload(
+            {
+                "type": "response.failed",
+                "response": {"error": {"message": "image blocked"}},
+            }
+        )
+
+
+def test_chatgpt_image_generation_handles_invalid_json_payloads():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "application/json"},
+        text="{not-json",
+    )
+
+    assert config._extract_image_payloads(raw_response) == []
+    assert config._get_parsed_payloads(raw_response) == [{}]
+
+
+def test_chatgpt_image_generation_ignores_non_dict_json_payload():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(status_code=200, json=[])
+
+    assert config._extract_image_payloads(raw_response) == []
+    assert config._get_parsed_payloads(raw_response) == []
+
+
+def test_chatgpt_image_generation_extracts_from_cyclic_nested_payload():
+    config = ChatGPTImageGenerationConfig()
+    payload = {
+        "type": "image_generation_call",
+        "result": "b64-result",
+        "b64_json": "b64-json",
+        "image": ["b64-image", 123],
+    }
+    payload["self"] = payload
+
+    assert config._extract_images_from_nested_value(payload) == [
+        "b64-result",
+        "b64-json",
+        "b64-image",
+    ]
+
+    cyclic_list = []
+    cyclic_list.append(cyclic_list)
+    assert config._extract_images_from_nested_value(cyclic_list) == []
+
+
+def test_chatgpt_image_generation_extracts_b64_from_streaming_completed_event():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "text/event-stream"},
+        text=(
+            'data: {"type":"image_generation.partial_image","b64_json":"partial-image"}\n\n'
+            'data: {"type":"image_generation.completed","b64_json":"final-image"}\n\n'
+            "data: [DONE]\n\n"
+        ),
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={"input": "draw a cat"},
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.data is not None
+    assert [item.b64_json for item in response.data] == ["final-image"]
+
+
+def test_chatgpt_image_generation_get_error_class():
+    config = ChatGPTImageGenerationConfig()
+
+    error = config.get_error_class(
+        error_message="bad request",
+        status_code=400,
+        headers={"x-request-id": "req-123"},
+    )
+
+    assert isinstance(error, OpenAIError)
+    assert error.status_code == 400
+    assert error.message == "bad request"

tests/test_litellm/llms/chatgpt/image_generation/test_usage.py

@@ -0,0 +1,199 @@
+import httpx
+import pytest
+
+from litellm.llms.chatgpt.image_generation import ChatGPTImageGenerationConfig
+from litellm.types.utils import ImageResponse
+from tests.test_litellm.llms.chatgpt.chatgpt_image_test_utils import mock_logging
+
+
+@pytest.fixture(autouse=True)
+def _chatgpt_token_dir(monkeypatch, tmp_path):
+    monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path))
+
+
+def test_chatgpt_image_generation_usage_helpers_ignore_invalid_payloads():
+    config = ChatGPTImageGenerationConfig()
+
+    assert config._get_image_generation_usage("not-a-dict") is None
+    assert config._get_image_generation_usage({"tool_usage": []}) is None
+    assert config._get_image_generation_usage({"tool_usage": {"image_gen": []}}) is None
+    assert (
+        config._get_image_generation_usage(
+            {"tool_usage": {"image_gen": {"input_tokens": 1}}}
+        )
+        is None
+    )
+    assert config._is_zero_image_usage(
+        {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
+    )
+
+
+def test_chatgpt_image_generation_extracts_tool_usage_from_completed_response():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "text/event-stream"},
+        text=(
+            'data: {"type":"response.completed","response":{"output":['
+            '{"type":"image_generation_call","result":"b64-image-data"}],'
+            '"usage":{"input_tokens":1732,"output_tokens":121,"total_tokens":1853},'
+            '"tool_usage":{"image_gen":{"input_tokens":108,'
+            '"input_tokens_details":{"image_tokens":0,"text_tokens":108},'
+            '"output_tokens":1756,'
+            '"output_tokens_details":{"image_tokens":1756,"text_tokens":0},'
+            '"total_tokens":1864}}}}\n\n'
+            "data: [DONE]\n\n"
+        ),
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={"input": "draw a cat"},
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.usage is not None
+    assert response.usage.input_tokens == 108
+    assert response.usage.input_tokens_details.text_tokens == 108
+    assert response.usage.input_tokens_details.image_tokens == 0
+    assert response.usage.output_tokens == 1756
+    assert response.usage.total_tokens == 1864
+
+
+def test_chatgpt_image_generation_prefers_completed_tool_usage():
+    config = ChatGPTImageGenerationConfig()
+    zero_usage = (
+        '"tool_usage":{"image_gen":{"input_tokens":0,'
+        '"input_tokens_details":{"image_tokens":0,"text_tokens":0},'
+        '"output_tokens":0,'
+        '"output_tokens_details":{"image_tokens":0,"text_tokens":0},'
+        '"total_tokens":0}}'
+    )
+    completed_usage = (
+        '"tool_usage":{"image_gen":{"input_tokens":105,'
+        '"input_tokens_details":{"image_tokens":0,"text_tokens":105},'
+        '"output_tokens":1372,'
+        '"output_tokens_details":{"image_tokens":1372,"text_tokens":0},'
+        '"total_tokens":1477}}'
+    )
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "text/event-stream"},
+        text=(
+            'data: {"type":"response.created","response":{'
+            f"{zero_usage}"
+            "}}\n\n"
+            'data: {"type":"response.in_progress","response":{'
+            f"{zero_usage}"
+            "}}\n\n"
+            'data: {"type":"response.image_generation_call.partial_image",'
+            '"partial_image_b64":"partial-image"}\n\n'
+            'data: {"type":"response.completed","response":{"output":['
+            '{"type":"image_generation_call","result":"b64-image-data"}],'
+            f"{completed_usage}"
+            ',"usage":{"input_tokens":2344,"output_tokens":118,"total_tokens":2462}}}\n\n'
+            "data: [DONE]\n\n"
+        ),
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={"input": "draw a cat"},
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.usage is not None
+    assert response.usage.input_tokens == 105
+    assert response.usage.input_tokens_details.text_tokens == 105
+    assert response.usage.input_tokens_details.image_tokens == 0
+    assert response.usage.output_tokens == 1372
+    assert response.usage.total_tokens == 1477
+
+
+def test_chatgpt_image_generation_extracts_usage_with_partial_image_payload():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "text/event-stream"},
+        text=(
+            "event: response.created\n"
+            'data: {"type":"response.created","response":{"tool_usage":{"image_gen":{'
+            '"input_tokens":0,"input_tokens_details":{"image_tokens":0,"text_tokens":0},'
+            '"output_tokens":0,"total_tokens":0}}}}\n\n'
+            "event: response.image_generation_call.partial_image\n"
+            'data: {"type":"response.image_generation_call.partial_image",'
+            '"partial_image_b64":"partial-image-data","size":"1536x1024"}\n\n'
+            "event: response.completed\n"
+            'data: {"type":"response.completed","response":{"output":[],'
+            '"tool_usage":{"image_gen":{"input_tokens":105,'
+            '"input_tokens_details":{"image_tokens":0,"text_tokens":105},'
+            '"output_tokens":1372,'
+            '"output_tokens_details":{"image_tokens":1372,"text_tokens":0},'
+            '"total_tokens":1477}},'
+            '"usage":{"input_tokens":2344,"output_tokens":118,"total_tokens":2462}}}\n\n'
+        ),
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={"input": "draw a cat"},
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.data is not None
+    assert response.data[0].b64_json == "partial-image-data"
+    assert response.usage is not None
+    assert response.usage.input_tokens == 105
+    assert response.usage.input_tokens_details.text_tokens == 105
+    assert response.usage.input_tokens_details.image_tokens == 0
+    assert response.usage.output_tokens == 1372
+    assert response.usage.total_tokens == 1477
+
+
+def test_chatgpt_image_generation_extracts_top_level_tool_usage():
+    config = ChatGPTImageGenerationConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        headers={"content-type": "text/event-stream"},
+        text=(
+            'data: {"type":"response.completed","response":{"output":['
+            '{"type":"image_generation_call","result":"b64-image-data"}]},'
+            '"tool_usage":{"image_gen":{"input_tokens":12,'
+            '"input_tokens_details":{"image_tokens":2,"text_tokens":10},'
+            '"output_tokens":34}}}\n\n'
+            "data: [DONE]\n\n"
+        ),
+    )
+
+    response = config.transform_image_generation_response(
+        model="gpt-image-2",
+        raw_response=raw_response,
+        model_response=ImageResponse(),
+        logging_obj=mock_logging(),
+        request_data={"input": "draw a cat"},
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+
+    assert response.usage is not None
+    assert response.usage.input_tokens == 12
+    assert response.usage.input_tokens_details.text_tokens == 10
+    assert response.usage.input_tokens_details.image_tokens == 2
+    assert response.usage.output_tokens == 34
+    assert response.usage.total_tokens == 46