diff --git a/litellm/llms/chatgpt/image_generation/transformation.py b/litellm/llms/chatgpt/image_generation/transformation.py index c5b27834e25..1c27464f807 100644 --- a/litellm/llms/chatgpt/image_generation/transformation.py +++ b/litellm/llms/chatgpt/image_generation/transformation.py @@ -120,7 +120,7 @@ class ChatGPTImageGenerationConfig(BaseImageGenerationConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - api_base = api_base or self.authenticator.get_api_base() or CHATGPT_API_BASE + api_base = self.authenticator.get_api_base() or CHATGPT_API_BASE api_base = self._canonicalize_codex_api_base(api_base) return f"{api_base}/responses" diff --git a/tests/image_gen_tests/test_chatgpt_image_generation.py b/tests/image_gen_tests/test_chatgpt_image_generation.py index c6e61fa5799..7cfffda5cb4 100644 --- a/tests/image_gen_tests/test_chatgpt_image_generation.py +++ b/tests/image_gen_tests/test_chatgpt_image_generation.py @@ -254,7 +254,7 @@ def test_chatgpt_image_generation_validate_environment_auth_error( @pytest.mark.parametrize( - "api_base, expected", + "server_api_base, expected", [ ( "https://chatgpt.com/backend-api", @@ -264,18 +264,27 @@ def test_chatgpt_image_generation_validate_environment_auth_error( "https://chatgpt.com/backend-api/responses", "https://chatgpt.com/backend-api/codex/responses", ), - ("https://example.test/custom/", "https://example.test/custom/responses"), + ( + "https://example.test/custom/", + "https://example.test/custom/responses", + ), ], ) -def test_chatgpt_image_generation_get_complete_url_canonicalizes_api_base( - monkeypatch, tmp_path, api_base, expected +def test_chatgpt_image_generation_get_complete_url_canonicalizes_server_api_base( + monkeypatch, tmp_path, server_api_base, expected ): monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) config = ChatGPTImageGenerationConfig() + class FakeAuthenticator: + def get_api_base(self): + return server_api_base + + config.authenticator = cast(Any, FakeAuthenticator()) + assert ( config.get_complete_url( - api_base=api_base, + api_base=None, api_key=None, model="gpt-image-2", optional_params={}, @@ -285,6 +294,30 @@ def test_chatgpt_image_generation_get_complete_url_canonicalizes_api_base( ) +def test_chatgpt_image_generation_get_complete_url_ignores_request_api_base( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + class FakeAuthenticator: + def get_api_base(self): + return "https://chatgpt.com/backend-api" + + config.authenticator = cast(Any, FakeAuthenticator()) + + assert ( + config.get_complete_url( + api_base="https://attacker.test/collect", + api_key=None, + model="gpt-image-2", + optional_params={}, + litellm_params={}, + ) + == "https://chatgpt.com/backend-api/codex/responses" + ) + + def test_chatgpt_image_generation_get_complete_url_uses_authenticator_api_base( monkeypatch, tmp_path ):