From a1c939b2ef74e72c74981d49e230b35957bbc4a9 Mon Sep 17 00:00:00 2001 From: Zero Clover Date: Wed, 25 Feb 2026 01:29:06 +0800 Subject: [PATCH 1/2] fix(image_generation): propagate extra_headers to OpenAI image generation Add headers parameter to image_generation() and aimage_generation() methods in OpenAI provider, and pass headers from images/main.py to ensure custom headers like cf-aig-authorization are properly forwarded to the OpenAI API. Aligns behavior with completion() method and Azure provider implementation. --- litellm/images/main.py | 1 + litellm/llms/openai/openai.py | 8 +++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 6c4c502a7b0..236266af6ad 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -483,6 +483,7 @@ def image_generation( # noqa: PLR0915 organization=organization, aimg_generation=aimg_generation, client=client, + headers=headers, ) elif custom_llm_provider == "bedrock": if model is None: diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index c7524925bd0..7020f796bb7 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1401,6 +1401,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=None, max_retries=None, organization: Optional[str] = None, + headers: Optional[dict] = None, ): response = None try: @@ -1414,6 +1415,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=client, ) + if headers: + data["extra_headers"] = headers response = await openai_aclient.images.generate(**data, timeout=timeout) # type: ignore stringified_response = response.model_dump() ## LOGGING @@ -1446,6 +1449,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=None, aimg_generation=None, organization: Optional[str] = None, + headers: Optional[dict] = None, ) -> ImageResponse: data = {} try: @@ -1455,7 +1459,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): raise OpenAIError(status_code=422, message="max retries must be an int") if aimg_generation is True: - return self.aimage_generation(data=data, prompt=prompt, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries, organization=organization) # type: ignore + return self.aimage_generation(data=data, prompt=prompt, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries, organization=organization, headers=headers) # type: ignore openai_client: OpenAI = self._get_openai_client( # type: ignore is_async=False, @@ -1480,6 +1484,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) ## COMPLETION CALL + if headers: + data["extra_headers"] = headers _response = openai_client.images.generate(**data, timeout=timeout) # type: ignore response = _response.model_dump() From a5b4dfe6bee59f49657695d872cd48f007625138 Mon Sep 17 00:00:00 2001 From: Zero Clover Date: Wed, 25 Feb 2026 01:51:18 +0800 Subject: [PATCH 2/2] test(image_generation): add tests for extra_headers propagation Verify that extra_headers are correctly forwarded to OpenAI's images.generate() in both sync and async paths, and that they are absent when not provided. --- ...t_openai_image_generation_extra_headers.py | 212 ++++++++++++++++++ 1 file changed, 212 insertions(+) create mode 100644 tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py 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 new file mode 100644 index 00000000000..33db9d33c1c --- /dev/null +++ b/tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -0,0 +1,212 @@ +""" +Unit tests for extra_headers propagation in OpenAI image generation. + +Verifies that extra_headers passed to litellm.image_generation() / +litellm.aimage_generation() are forwarded to the OpenAI API client as +extra_headers in the images.generate() call. +""" + +import os +import sys +from unittest.mock import MagicMock, AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +@pytest.fixture +def openai_chat_completions(): + return OpenAIChatCompletion() + + +@pytest.fixture +def mock_logging_obj(): + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.post_call = MagicMock() + return 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 + ): + """Sync image_generation should pass headers as extra_headers to images.generate().""" + 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" + + test_headers = {"cf-aig-authorization": "Bearer custom-token"} + + openai_chat_completions.image_generation( + model="dall-e-3", + prompt="A white cat", + timeout=60.0, + optional_params={}, + logging_obj=mock_logging_obj, + api_key="test-key", + headers=test_headers, + client=mock_openai_client, + ) + + _, 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 + ): + """Sync image_generation without headers should not inject extra_headers.""" + 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="dall-e-3", + prompt="A white cat", + timeout=60.0, + optional_params={}, + logging_obj=mock_logging_obj, + api_key="test-key", + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + 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 aimage_generation should pass headers as extra_headers to images.generate().""" + 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 = AsyncMock(return_value=mock_image_data) + mock_openai_client.api_key = "test-key" + + test_headers = {"cf-aig-authorization": "Bearer custom-token"} + + await openai_chat_completions.aimage_generation( + prompt="A white cat", + data={"model": "dall-e-3", "prompt": "A white cat"}, + model_response=MagicMock(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + headers=test_headers, + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + 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 aimage_generation without headers should not inject extra_headers.""" + 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 = AsyncMock(return_value=mock_image_data) + mock_openai_client.api_key = "test-key" + + await openai_chat_completions.aimage_generation( + prompt="A white cat", + data={"model": "dall-e-3", "prompt": "A white cat"}, + model_response=MagicMock(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + client=mock_openai_client, + ) + + _, kwargs = mock_openai_client.images.generate.call_args + assert "extra_headers" not in kwargs + + 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: + mock_aimage_gen.return_value = MagicMock() + + test_headers = {"x-custom-header": "value"} + + openai_chat_completions.image_generation( + model="dall-e-3", + prompt="A white cat", + timeout=60.0, + optional_params={}, + logging_obj=mock_logging_obj, + api_key="test-key", + aimg_generation=True, + headers=test_headers, + ) + + mock_aimage_gen.assert_called_once() + call_kwargs = mock_aimage_gen.call_args[1] + assert call_kwargs["headers"] == test_headers + + +class TestImageGenerationEntryPointHeaders: + """Test that litellm.image_generation() passes headers through to the OpenAI provider.""" + + @pytest.mark.asyncio + async def test_extra_headers_reach_openai_provider(self): + """End-to-end: extra_headers from litellm.aimage_generation() reach OpenAI images.generate().""" + import litellm + + 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 = AsyncMock(return_value=mock_image_data) + mock_openai_client.api_key = "test-key" + mock_openai_client._base_url._uri_reference = "https://api.openai.com" + + test_headers = {"cf-aig-authorization": "Bearer my-secret"} + + await litellm.aimage_generation( + model="dall-e-3", + prompt="A white cat", + extra_headers=test_headers, + client=mock_openai_client, + api_key="test-key", + ) + + mock_openai_client.images.generate.assert_called_once() + _, kwargs = mock_openai_client.images.generate.call_args + assert kwargs.get("extra_headers") == test_headers