From ed72e3200bc0ca8438822bc13e1787e06d86e22a Mon Sep 17 00:00:00 2001 From: CrystalVibe28 <82704512+CrystalVibe28@users.noreply.github.com> Date: Sat, 16 May 2026 19:38:27 +0800 Subject: [PATCH] refactor(chatgpt): split image transformation modules --- litellm/__init__.py | 4 +- litellm/_lazy_imports_registry.py | 4 +- litellm/llms/chatgpt/image_edit/__init__.py | 3 + .../llms/chatgpt/image_edit/transformation.py | 179 +++ .../llms/chatgpt/image_generation/__init__.py | 4 +- .../generation_transformation.py | 319 +++++ .../image_generation/response_parsing.py | 235 ++++ .../image_generation/transformation.py | 668 ----------- litellm/utils.py | 2 +- .../llms/chatgpt/chatgpt_image_test_utils.py | 10 + .../chatgpt/image_edit/test_transformation.py | 189 +++ .../test_request_transformation.py | 395 +++++++ .../image_generation/test_response_parsing.py | 211 ++++ .../chatgpt/image_generation/test_usage.py | 199 ++++ .../chatgpt/test_chatgpt_image_generation.py | 1040 ----------------- 15 files changed, 1748 insertions(+), 1714 deletions(-) create mode 100644 litellm/llms/chatgpt/image_edit/__init__.py create mode 100644 litellm/llms/chatgpt/image_edit/transformation.py create mode 100644 litellm/llms/chatgpt/image_generation/generation_transformation.py create mode 100644 litellm/llms/chatgpt/image_generation/response_parsing.py delete mode 100644 litellm/llms/chatgpt/image_generation/transformation.py create mode 100644 tests/test_litellm/llms/chatgpt/chatgpt_image_test_utils.py create mode 100644 tests/test_litellm/llms/chatgpt/image_edit/test_transformation.py create mode 100644 tests/test_litellm/llms/chatgpt/image_generation/test_request_transformation.py create mode 100644 tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py create mode 100644 tests/test_litellm/llms/chatgpt/image_generation/test_usage.py delete mode 100644 tests/test_litellm/llms/chatgpt/test_chatgpt_image_generation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 8f6cb0fd137..be0f6dc14f3 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 4305fdf80d7..54b2684ec06 100644 --- a/litellm/_lazy_imports_registry.py +++ b/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"), diff --git a/litellm/llms/chatgpt/image_edit/__init__.py b/litellm/llms/chatgpt/image_edit/__init__.py new file mode 100644 index 00000000000..427c840f70c --- /dev/null +++ b/litellm/llms/chatgpt/image_edit/__init__.py @@ -0,0 +1,3 @@ +from .transformation import ChatGPTImageEditConfig + +__all__ = ["ChatGPTImageEditConfig"] diff --git a/litellm/llms/chatgpt/image_edit/transformation.py b/litellm/llms/chatgpt/image_edit/transformation.py new file mode 100644 index 00000000000..e34bb2ad81d --- /dev/null +++ b/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, + ) diff --git a/litellm/llms/chatgpt/image_generation/__init__.py b/litellm/llms/chatgpt/image_generation/__init__.py index 041ed1611e5..ebf7e0d7cd3 100644 --- a/litellm/llms/chatgpt/image_generation/__init__.py +++ b/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"] diff --git a/litellm/llms/chatgpt/image_generation/generation_transformation.py b/litellm/llms/chatgpt/image_generation/generation_transformation.py new file mode 100644 index 00000000000..9bc6f91ce20 --- /dev/null +++ b/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, + ) diff --git a/litellm/llms/chatgpt/image_generation/response_parsing.py b/litellm/llms/chatgpt/image_generation/response_parsing.py new file mode 100644 index 00000000000..2b3b8545351 --- /dev/null +++ b/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 diff --git a/litellm/llms/chatgpt/image_generation/transformation.py b/litellm/llms/chatgpt/image_generation/transformation.py deleted file mode 100644 index 74a85ee3a47..00000000000 --- a/litellm/llms/chatgpt/image_generation/transformation.py +++ /dev/null @@ -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, - ) diff --git a/litellm/utils.py b/litellm/utils.py index 8074c6c4ea0..da79e54529d 100644 --- a/litellm/utils.py +++ b/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: diff --git a/tests/test_litellm/llms/chatgpt/chatgpt_image_test_utils.py b/tests/test_litellm/llms/chatgpt/chatgpt_image_test_utils.py new file mode 100644 index 00000000000..23b080883b3 --- /dev/null +++ b/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() diff --git a/tests/test_litellm/llms/chatgpt/image_edit/test_transformation.py b/tests/test_litellm/llms/chatgpt/image_edit/test_transformation.py new file mode 100644 index 00000000000..0fae768c907 --- /dev/null +++ b/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())) diff --git a/tests/test_litellm/llms/chatgpt/image_generation/test_request_transformation.py b/tests/test_litellm/llms/chatgpt/image_generation/test_request_transformation.py new file mode 100644 index 00000000000..b0d0a03b88f --- /dev/null +++ b/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"} diff --git a/tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py b/tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py new file mode 100644 index 00000000000..d8851f345eb --- /dev/null +++ b/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" diff --git a/tests/test_litellm/llms/chatgpt/image_generation/test_usage.py b/tests/test_litellm/llms/chatgpt/image_generation/test_usage.py new file mode 100644 index 00000000000..703b10b3dd7 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_image_generation.py b/tests/test_litellm/llms/chatgpt/test_chatgpt_image_generation.py deleted file mode 100644 index 8535fa07729..00000000000 --- a/tests/test_litellm/llms/chatgpt/test_chatgpt_image_generation.py +++ /dev/null @@ -1,1040 +0,0 @@ -from io import BytesIO -from typing import Any, cast - -import httpx -import pytest - -from litellm.exceptions import AuthenticationError -from litellm.llms.chatgpt.common_utils import GetAccessTokenError -from litellm.llms.chatgpt.image_generation.transformation import ( - ChatGPTImageEditConfig, - ChatGPTImageGenerationConfig, -) -from litellm.llms.openai.common_utils import OpenAIError -from litellm.types.utils import LlmProviders -from litellm.types.utils import ImageResponse -from litellm.utils import ProviderConfigManager - - -class MockLogging: - def post_call(self, *args, **kwargs): - pass - - -def mock_logging() -> Any: - return MockLogging() - - -def test_chatgpt_image_generation_transforms_request(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path, optional_params, error -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - config = ProviderConfigManager.get_provider_image_generation_config( - model="gpt-image-2", - provider=LlmProviders.CHATGPT, - ) - - assert isinstance(config, ChatGPTImageGenerationConfig) - - -def test_chatgpt_image_edit_config_registered(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - config = ChatGPTImageEditConfig() - - assert config.use_multipart_form_data() is False - - -def test_chatgpt_image_edit_requires_image(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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_generation_only_supports_prompt_output_format_and_size( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - config = ChatGPTImageGenerationConfig() - - assert config.get_supported_openai_params("gpt-image-2") == [ - "output_format", - "size", - ] - - -def test_chatgpt_image_generation_rejects_unsupported_optional_params( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path, server_api_base, expected -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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_extracts_b64_from_sse_completed_response( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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_uses_optional_responses_model(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path, model, optional_params, error -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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_extracts_b64_from_deep_nested_payload( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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_usage_helpers_ignore_invalid_payloads( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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 - - -def test_chatgpt_image_generation_extracts_b64_from_streaming_completed_event( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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" - - -def test_chatgpt_image_generation_map_openai_params_keeps_existing_value( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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"} - - -def test_chatgpt_image_edit_delegates_environment_and_url(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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(monkeypatch, tmp_path): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - 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( - monkeypatch, tmp_path -): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(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()))