fix: forward OpenAI organization for image generation (#16607)

litellm/images/main.py

@@ -400,6 +400,8 @@ def image_generation(  # noqa: PLR0915
             or custom_llm_provider == LlmProviders.LITELLM_PROXY.value
             or custom_llm_provider in litellm.openai_compatible_providers
         ):
+            # Forward OpenAI organization if present (set by proxy pre-call utils)
+            organization: Optional[str] = kwargs.get("organization", None)
             model_response = openai_chat_completions.image_generation(
                 model=model,
                 prompt=prompt,
@@ -409,6 +411,7 @@ def image_generation(  # noqa: PLR0915
                 logging_obj=litellm_logging_obj,
                 optional_params=optional_params,
                 model_response=model_response,
+                organization=organization,
                 aimg_generation=aimg_generation,
                 client=client,
             )

litellm/llms/openai/openai.py

@@ -1285,6 +1285,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         api_base: Optional[str] = None,
         client=None,
         max_retries=None,
+        organization: Optional[str] = None,
     ):
         response = None
         try:
@@ -1294,6 +1295,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                 api_base=api_base,
                 timeout=timeout,
                 max_retries=max_retries,
+                organization=organization,
                 client=client,
             )
 
@@ -1328,6 +1330,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
         model_response: Optional[ImageResponse] = None,
         client=None,
         aimg_generation=None,
+        organization: Optional[str] = None,
     ) -> ImageResponse:
         data = {}
         try:
@@ -1337,7 +1340,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)  # 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)  # type: ignore
 
             openai_client: OpenAI = self._get_openai_client(  # type: ignore
                 is_async=False,
@@ -1345,6 +1348,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
                 api_base=api_base,
                 timeout=timeout,
                 max_retries=max_retries,
+                organization=organization,
                 client=client,
             )
 

tests/llm_translation/test_openai.py

@@ -337,6 +337,48 @@ def test_openai_max_retries_0(mock_get_openai_client):
     assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0
 
 
+@patch("litellm.main.openai_chat_completions._get_openai_client")
+def test_openai_image_generation_forwards_organization(mock_get_openai_client):
+    """Ensure organization flows to OpenAI client for image generation."""
+
+    class _DummyImages:
+        def generate(self, **kwargs):  # type: ignore
+            class _Resp:
+                def model_dump(self_inner):  # minimal OpenAI ImagesResponse shape
+                    return {
+                        "created": 123,
+                        "data": [{"url": "http://example.com/image.png"}],
+                        "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
+                    }
+
+            return _Resp()
+
+    class _DummyClient:
+        def __init__(self):
+            self.api_key = "sk-test"
+
+            class _BaseURL:
+                _uri_reference = "https://api.openai.com/v1"
+
+            self._base_url = _BaseURL()
+            self.images = _DummyImages()
+
+    mock_get_openai_client.return_value = _DummyClient()
+
+    org = "org_test_123"
+    resp = litellm.image_generation(
+        model="gpt-image-1",
+        prompt="A cute baby sea otter",
+        organization=org,
+    )
+
+    # Assert organization forwarded into OpenAI client factory
+    assert mock_get_openai_client.call_args.kwargs.get("organization") == org
+
+    # Basic sanity on response shape
+    assert hasattr(resp, "data") and len(resp.data) == 1
+
+
 @pytest.mark.parametrize("model", ["o1", "o3-mini"])
 def test_o1_parallel_tool_calls(model):
     litellm.completion(