Merge a3da59f2be into 10b3deae0d

litellm/llms/openai/openai.py

@@ -377,6 +377,12 @@ _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({})
 _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body"))
 
 
+def _as_mapping(value: object) -> Mapping[str, object]:
+    if not isinstance(value, dict):
+        return types.MappingProxyType({})
+    return types.MappingProxyType({key: item for key, item in value.items()})
+
+
 def _raw_chat_response(
     result: ChatCompletion | Stream[ChatCompletionChunk] | AsyncStream[ChatCompletionChunk],
 ) -> LegacyAPIResponse[ChatCompletion]:
@@ -398,6 +404,41 @@ def _request_without_sdk_transform(
     return body, options
 
 
+def _image_generation_logged_body(data: Mapping[str, object]) -> dict[str, object]:
+    extra_body: Final = {
+        key: item for key, item in _as_mapping(data.get("extra_body")).items() if key != "extra_headers"
+    }
+    extra_query: Final = dict(_as_mapping(data.get("extra_query")))
+    return {  # mutable-ok: loggers isinstance-check this payload as a dict
+        **{key: item for key, item in data.items() if key not in _SDK_OPTION_KEYS},
+        **({} if not extra_query else {"extra_query": extra_query}),
+        **({} if not extra_body else {"extra_body": extra_body}),
+    }
+
+
+def _image_generation_request_data(
+    data: Mapping[str, object], headers: Mapping[str, object] | None
+) -> tuple[dict[str, object], Mapping[str, object]]:
+    extra_body_source: Final = _as_mapping(data.get("extra_body"))
+    extra_headers: Final = types.MappingProxyType(
+        {
+            **_as_mapping(extra_body_source.get("extra_headers")),
+            **_as_mapping(data.get("extra_headers")),
+            **_as_mapping(headers),
+        }
+    )
+    logged_body: Final = _image_generation_logged_body(data)
+    request_data: Final = types.MappingProxyType(
+        dict(
+            (
+                *logged_body.items(),
+                *((("extra_headers", extra_headers),) if extra_headers else ()),
+            )
+        )
+    )
+    return logged_body, request_data
+
+
 class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
     def __init__(self) -> None:
         super().__init__()
@@ -1547,6 +1588,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                 client=client,
             )
 
+            logged_body, request_data = _image_generation_request_data(data, headers)
             logging_obj.pre_call(
                 input=prompt,
                 api_key=openai_aclient.api_key,
@@ -1554,11 +1596,10 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                     "headers": {"Authorization": f"Bearer {openai_aclient.api_key}"},
                     "api_base": str(openai_aclient.base_url),
                     "acompletion": True,
-                    "complete_input_dict": data,
+                    "complete_input_dict": logged_body,
                 },
             )
 
-            request_data: Final = {**data, "extra_headers": headers} if headers else data
             raw_response: Final = await openai_aclient.images.with_raw_response.generate(
                 **request_data, timeout=timeout
             )
@@ -1567,7 +1608,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             logging_obj.post_call(
                 input=prompt,
                 api_key=api_key,
-                additional_args={"complete_input_dict": data},
+                additional_args={"complete_input_dict": logged_body},
                 original_response=stringified_response,
             )
             image_response: Final[ImageResponse] = convert_to_model_response_object(
@@ -1582,6 +1623,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             logging_obj.post_call(
                 input=prompt,
                 api_key=api_key,
+                additional_args={"complete_input_dict": _image_generation_logged_body(data)},
                 original_response=str(e),
             )
             raise e
@@ -1636,6 +1678,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                 client=client,
             )
 
+            logged_body, request_data = _image_generation_request_data(data, headers)
+
             ## LOGGING
             logging_obj.pre_call(
                 input=prompt,
@@ -1644,12 +1688,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                     "headers": {"Authorization": f"Bearer {openai_client.api_key}"},
                     "api_base": openai_client._base_url._uri_reference,  # pyright: ignore[reportPrivateUsage]  # SDK URL internals
                     "acompletion": True,
-                    "complete_input_dict": data,
+                    "complete_input_dict": logged_body,
                 },
             )
 
             ## COMPLETION CALL
-            request_data: Final = {**data, "extra_headers": headers} if headers else data
             raw_response: Final = openai_client.images.with_raw_response.generate(**request_data, timeout=timeout)
 
             response: Final = raw_response.parse().model_dump()
@@ -1657,7 +1700,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             logging_obj.post_call(
                 input=prompt,
                 api_key=api_key,
-                additional_args={"complete_input_dict": data},
+                additional_args={"complete_input_dict": logged_body},
                 original_response=response,
             )
             image_response: Final[ImageResponse] = convert_to_model_response_object(
@@ -1672,7 +1715,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             logging_obj.post_call(
                 input=prompt,
                 api_key=api_key,
-                additional_args={"complete_input_dict": data},
+                additional_args={"complete_input_dict": _image_generation_logged_body(data)},
                 original_response=str(e),
             )
             raise e
@@ -1681,7 +1724,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
             logging_obj.post_call(
                 input=prompt,
                 api_key=api_key,
-                additional_args={"complete_input_dict": data},
+                additional_args={"complete_input_dict": _image_generation_logged_body(data)},
                 original_response=str(e),
             )
             if hasattr(e, "status_code"):

tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py

@@ -6,11 +6,11 @@ litellm.aimage_generation() are forwarded to the OpenAI API client as
 extra_headers in the images.generate() call.
 """
 
-from unittest.mock import MagicMock, AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-
+from litellm.llms.openai.common_utils import OpenAIError
 from litellm.llms.openai.openai import OpenAIChatCompletion
 from litellm.types.utils import ImageResponse
 
@@ -21,6 +21,21 @@ def _raw_image_response(mock_image_data):
     raw_response.headers = {"x-request-id": "req-image"}
     return raw_response
 
+_SECRET_HEADERS = {"Authorization": "Bearer leaked-token", "cf-aig-authorization": "Bearer custom-token"}
+
+
+def _assert_logged_body_omits_credentials(logged_body: object) -> None:
+    assert isinstance(logged_body, dict)
+    assert "extra_headers" not in logged_body
+    assert "Authorization" not in logged_body
+    assert "cf-aig-authorization" not in logged_body
+    extra_body = logged_body.get("extra_body") or {}
+    assert "extra_headers" not in extra_body
+    assert "Bearer leaked-token" not in str(logged_body)
+    assert extra_body.get("foo") == "bar"
+    assert logged_body.get("extra_query") == {"user": "alice"}
+    assert logged_body.get("size") == "1024x1024"
+
 
 @pytest.fixture
 def openai_chat_completions():
@@ -38,9 +53,7 @@ def mock_logging_obj():
 class TestImageGenerationExtraHeaders:
     """Test that extra_headers are properly injected into OpenAI image generation calls."""
 
-    def test_sync_image_generation_with_headers(
-        self, openai_chat_completions, mock_logging_obj
-    ):
+    def test_sync_image_generation_with_headers(self, openai_chat_completions, mock_logging_obj):
         """Sync image_generation should pass headers as extra_headers to images.generate()."""
         mock_image_data = MagicMock()
         mock_image_data.model_dump.return_value = {
@@ -69,9 +82,7 @@ class TestImageGenerationExtraHeaders:
         _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args
         assert kwargs.get("extra_headers") == test_headers
 
-    def test_sync_image_generation_without_headers(
-        self, openai_chat_completions, mock_logging_obj
-    ):
+    def test_sync_image_generation_without_headers(self, openai_chat_completions, mock_logging_obj):
         """Sync image_generation without headers should not inject extra_headers."""
         mock_image_data = MagicMock()
         mock_image_data.model_dump.return_value = {
@@ -98,9 +109,7 @@ class TestImageGenerationExtraHeaders:
         assert "extra_headers" not in kwargs
 
     @pytest.mark.asyncio
-    async def test_async_image_generation_with_headers(
-        self, openai_chat_completions, mock_logging_obj
-    ):
+    async def test_async_image_generation_with_headers(self, openai_chat_completions, mock_logging_obj):
         """Async aimage_generation should pass headers as extra_headers to images.generate()."""
         mock_image_data = MagicMock()
         mock_image_data.model_dump.return_value = {
@@ -131,9 +140,7 @@ class TestImageGenerationExtraHeaders:
         assert kwargs.get("extra_headers") == test_headers
 
     @pytest.mark.asyncio
-    async def test_async_image_generation_without_headers(
-        self, openai_chat_completions, mock_logging_obj
-    ):
+    async def test_async_image_generation_without_headers(self, openai_chat_completions, mock_logging_obj):
         """Async aimage_generation without headers should not inject extra_headers."""
         mock_image_data = MagicMock()
         mock_image_data.model_dump.return_value = {
@@ -207,20 +214,14 @@ class TestImageGenerationExtraHeaders:
                 client=mock_openai_client,
             )
 
-        logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"][
-            "complete_input_dict"
-        ]
+        logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"]["complete_input_dict"]
         assert "extra_headers" not in logged_body
         _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args
         assert kwargs.get("extra_headers") == test_headers
 
-    def test_sync_image_generation_forwards_headers_to_async(
-        self, openai_chat_completions, mock_logging_obj
-    ):
+    def test_sync_image_generation_forwards_headers_to_async(self, openai_chat_completions, mock_logging_obj):
         """When aimg_generation=True, image_generation should forward headers to aimage_generation."""
-        with patch.object(
-            openai_chat_completions, "aimage_generation"
-        ) as mock_aimage_gen:
+        with patch.object(openai_chat_completions, "aimage_generation") as mock_aimage_gen:
             mock_aimage_gen.return_value = MagicMock()
 
             test_headers = {"x-custom-header": "value"}
@@ -240,6 +241,157 @@ class TestImageGenerationExtraHeaders:
             call_kwargs = mock_aimage_gen.call_args[1]
             assert call_kwargs["headers"] == test_headers
 
+    @pytest.mark.parametrize("is_async", [False, True])
+    @pytest.mark.asyncio
+    async def test_optional_params_extra_headers_never_enter_the_json_body(
+        self, openai_chat_completions, mock_logging_obj, is_async
+    ):
+        """Proxy credentials put extra_headers on optional_params (and often
+        extra_body). Those must reach images.generate only as the SDK
+        extra_headers kwarg. OpenAI rejects extra_headers in the JSON body."""
+        mock_image_data = MagicMock()
+        mock_image_data.model_dump.return_value = {
+            "created": 1700000000,
+            "data": [{"url": "https://example.com/image.png"}],
+        }
+
+        mock_openai_client = MagicMock()
+        mock_openai_client.api_key = "test-key"
+        mock_openai_client._base_url._uri_reference = "https://api.openai.com"
+
+        test_headers = {"cf-aig-authorization": "Bearer custom-token"}
+        optional_params = {
+            "size": "1024x1024",
+            "quality": "low",
+            "extra_headers": test_headers,
+            "extra_body": {"extra_headers": test_headers, "foo": "bar"},
+            "extra_query": {"user": "alice"},
+        }
+
+        if is_async:
+            mock_openai_client.images.with_raw_response.generate = AsyncMock(
+                return_value=_raw_image_response(mock_image_data)
+            )
+            await openai_chat_completions.aimage_generation(
+                prompt="A white cat",
+                data={"model": "gpt-image-2", "prompt": "A white cat", **optional_params},
+                model_response=ImageResponse(),
+                timeout=60.0,
+                logging_obj=mock_logging_obj,
+                api_key="test-key",
+                client=mock_openai_client,
+            )
+        else:
+            mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data)
+            openai_chat_completions.image_generation(
+                model="gpt-image-2",
+                prompt="A white cat",
+                timeout=60.0,
+                optional_params=optional_params,
+                logging_obj=mock_logging_obj,
+                api_key="test-key",
+                client=mock_openai_client,
+            )
+
+        _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args
+        assert kwargs.get("extra_headers") == test_headers
+        assert kwargs.get("size") == "1024x1024"
+        assert kwargs.get("quality") == "low"
+        extra_body = kwargs.get("extra_body") or {}
+        assert "extra_headers" not in extra_body
+        assert extra_body.get("foo") == "bar"
+        assert kwargs.get("extra_query") == {"user": "alice"}
+
+        logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"]["complete_input_dict"]
+        _assert_logged_body_omits_credentials(logged_body)
+
+    def test_optional_params_headers_merge_prefers_caller_headers(self, openai_chat_completions, mock_logging_obj):
+        """When both optional_params and the headers arg carry extra_headers,
+        the explicit headers arg wins, matching the previous overwrite."""
+        mock_image_data = MagicMock()
+        mock_image_data.model_dump.return_value = {
+            "created": 1700000000,
+            "data": [{"url": "https://example.com/image.png"}],
+        }
+
+        mock_openai_client = MagicMock()
+        mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data)
+        mock_openai_client.api_key = "test-key"
+        mock_openai_client._base_url._uri_reference = "https://api.openai.com"
+
+        openai_chat_completions.image_generation(
+            model="gpt-image-2",
+            prompt="A white cat",
+            timeout=60.0,
+            optional_params={
+                "extra_headers": {"cf-aig-authorization": "Bearer from-optional"},
+            },
+            logging_obj=mock_logging_obj,
+            api_key="test-key",
+            headers={"cf-aig-authorization": "Bearer from-headers"},
+            client=mock_openai_client,
+        )
+
+        _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args
+        assert kwargs.get("extra_headers") == {"cf-aig-authorization": "Bearer from-headers"}
+
+    @pytest.mark.parametrize(
+        "is_async,raised",
+        [
+            (False, OpenAIError(status_code=400, message="bad image params")),
+            (False, RuntimeError("forced failure")),
+            (True, OpenAIError(status_code=400, message="bad image params")),
+            (True, RuntimeError("forced failure")),
+        ],
+        ids=["sync-openai-error", "sync-runtime-error", "async-openai-error", "async-runtime-error"],
+    )
+    @pytest.mark.asyncio
+    async def test_failure_telemetry_omits_extra_headers(
+        self, openai_chat_completions, mock_logging_obj, is_async, raised
+    ):
+        """post_call on a generate() failure must use the sanitized body, not raw data."""
+        mock_openai_client = MagicMock()
+        mock_openai_client.api_key = "test-key"
+        mock_openai_client._base_url._uri_reference = "https://api.openai.com"
+
+        optional_params = {
+            "size": "1024x1024",
+            "extra_headers": _SECRET_HEADERS,
+            "extra_body": {"extra_headers": _SECRET_HEADERS, "foo": "bar"},
+            "extra_query": {"user": "alice"},
+        }
+
+        if is_async:
+            mock_openai_client.images.with_raw_response.generate = AsyncMock(side_effect=raised)
+            with pytest.raises(type(raised)):
+                await openai_chat_completions.aimage_generation(
+                    prompt="A white cat",
+                    data={"model": "gpt-image-2", "prompt": "A white cat", **optional_params},
+                    model_response=MagicMock(),
+                    timeout=60.0,
+                    logging_obj=mock_logging_obj,
+                    api_key="test-key",
+                    headers=_SECRET_HEADERS,
+                    client=mock_openai_client,
+                )
+        else:
+            mock_openai_client.images.with_raw_response.generate.side_effect = raised
+            with pytest.raises(OpenAIError):
+                openai_chat_completions.image_generation(
+                    model="gpt-image-2",
+                    prompt="A white cat",
+                    timeout=60.0,
+                    optional_params=optional_params,
+                    logging_obj=mock_logging_obj,
+                    api_key="test-key",
+                    headers=_SECRET_HEADERS,
+                    client=mock_openai_client,
+                )
+
+        mock_openai_client.images.with_raw_response.generate.assert_called_once()
+        logged_body = mock_logging_obj.post_call.call_args[1]["additional_args"]["complete_input_dict"]
+        _assert_logged_body_omits_credentials(logged_body)
+
 
 class TestImageGenerationEntryPointHeaders:
     """Test that litellm.image_generation() passes headers through to the OpenAI provider."""
@@ -275,3 +427,4 @@ class TestImageGenerationEntryPointHeaders:
         mock_openai_client.images.with_raw_response.generate.assert_called_once()
         _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args
         assert kwargs.get("extra_headers") == test_headers
+        assert "extra_headers" not in (kwargs.get("extra_body") or {})