From f1c2e169dedd3bbee8fe9ea4dfaf032b9d653eb3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 11 Sep 2026 02:04:01 +0000 Subject: [PATCH] fix(openai): keep extra_headers off the image generation JSON body Proxy credentials put extra_headers on optional_params and extra_body. Pop them out of the Images request body and pass them only as the SDK header kwarg so OpenAI no longer returns Unknown parameter extra_headers. Co-authored-by: Zsanz3 --- litellm/llms/openai/openai.py | 54 ++++++-- ...t_openai_image_generation_extra_headers.py | 123 ++++++++++++++---- 2 files changed, 144 insertions(+), 33 deletions(-) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 5f04ebe0c01..0c6cff5bca5 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -340,6 +340,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 _embedding_request_without_sdk_defaults( data: Mapping[str, object], timeout: float | httpx.Timeout ) -> tuple[Mapping[str, object], RequestOptions]: @@ -356,6 +362,37 @@ def _embedding_request_without_sdk_defaults( return body, options +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_body: Final = types.MappingProxyType( + {key: item for key, item in extra_body_source.items() if key != "extra_headers"} + ) + extra_headers: Final = types.MappingProxyType( + { + **_as_mapping(extra_body_source.get("extra_headers")), + **_as_mapping(data.get("extra_headers")), + **_as_mapping(headers), + } + ) + extra_query: Final = _as_mapping(data.get("extra_query")) + logged_body: Final = { # 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 + } + request_data: Final = types.MappingProxyType( + dict( + ( + *logged_body.items(), + *((("extra_headers", extra_headers),) if extra_headers else ()), + *((("extra_query", extra_query),) if extra_query else ()), + *((("extra_body", extra_body),) if extra_body else ()), + ) + ) + ) + return logged_body, request_data + + class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): def __init__(self) -> None: super().__init__() @@ -1415,6 +1452,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, @@ -1422,20 +1460,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): "headers": {"Authorization": f"Bearer {openai_aclient.api_key}"}, # mutable-ok: logged header map "api_base": str(openai_aclient.base_url), "acompletion": True, - "complete_input_dict": data, + "complete_input_dict": logged_body, }, ) - request_data: Final = ( # mutable-ok: the OpenAI SDK takes the request body as a dict - {**data, "extra_headers": headers} if headers else data - ) response = await openai_aclient.images.generate(**request_data, timeout=timeout) stringified_response: Final = response.model_dump() ## LOGGING 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, ) return convert_to_model_response_object( @@ -1499,6 +1534,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=client, ) + logged_body, request_data = _image_generation_request_data(data, headers) + ## LOGGING logging_obj.pre_call( input=prompt, @@ -1507,14 +1544,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): "headers": {"Authorization": f"Bearer {openai_client.api_key}"}, "api_base": openai_client._base_url._uri_reference, "acompletion": True, - "complete_input_dict": data, + "complete_input_dict": logged_body, }, ) ## COMPLETION CALL - request_data: Final = ( # mutable-ok: the OpenAI SDK takes the request body as a dict - {**data, "extra_headers": headers} if headers else data - ) _response: Final = openai_client.images.generate(**request_data, timeout=timeout) response: Final = _response.model_dump() @@ -1522,7 +1556,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, ) return convert_to_model_response_object( diff --git a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index 55ef74abd7b..b422e62b9c2 100644 --- a/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -6,11 +6,10 @@ 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.openai import OpenAIChatCompletion @@ -30,9 +29,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 = { @@ -61,9 +58,7 @@ class TestImageGenerationExtraHeaders: _, kwargs = mock_openai_client.images.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 = { @@ -90,9 +85,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 = { @@ -121,9 +114,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 = { @@ -193,20 +184,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.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"} @@ -226,6 +211,97 @@ 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"}, + } + + if is_async: + mock_openai_client.images.generate = AsyncMock(return_value=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=MagicMock(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + client=mock_openai_client, + ) + else: + mock_openai_client.images.generate.return_value = 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.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" + + logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"]["complete_input_dict"] + assert "extra_headers" not in logged_body + assert "extra_headers" not in (logged_body.get("extra_body") or {}) + + 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.generate.return_value = 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.generate.call_args + assert kwargs.get("extra_headers") == {"cf-aig-authorization": "Bearer from-headers"} + class TestImageGenerationEntryPointHeaders: """Test that litellm.image_generation() passes headers through to the OpenAI provider.""" @@ -259,3 +335,4 @@ class TestImageGenerationEntryPointHeaders: mock_openai_client.images.generate.assert_called_once() _, kwargs = mock_openai_client.images.generate.call_args assert kwargs.get("extra_headers") == test_headers + assert "extra_headers" not in (kwargs.get("extra_body") or {})