diff --git a/litellm/images/main.py b/litellm/images/main.py index ce6da640f9..333a751b04 100644 --- a/litellm/images/main.py +++ b/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, ) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 492ed62423..2949e35e5e 100644 --- a/litellm/llms/openai/openai.py +++ b/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, ) diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 1b96a99621..ba1d9e6ac2 100644 --- a/tests/llm_translation/test_openai.py +++ b/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(