diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6e22ef9112e..efe6846b750 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -377,6 +377,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 _raw_chat_response( result: ChatCompletion | Stream[ChatCompletionChunk] | AsyncStream[ChatCompletionChunk], ) -> LegacyAPIResponse[ChatCompletion]: @@ -398,6 +404,41 @@ def _request_without_sdk_transform( return body, options +def _image_generation_logged_body(data: Mapping[str, object]) -> dict[str, object]: + extra_body: Final = { + key: item for key, item in _as_mapping(data.get("extra_body")).items() if key != "extra_headers" + } + extra_query: Final = dict(_as_mapping(data.get("extra_query"))) + return { # 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}, + **({} if not extra_query else {"extra_query": extra_query}), + **({} if not extra_body else {"extra_body": extra_body}), + } + + +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_headers: Final = types.MappingProxyType( + { + **_as_mapping(extra_body_source.get("extra_headers")), + **_as_mapping(data.get("extra_headers")), + **_as_mapping(headers), + } + ) + logged_body: Final = _image_generation_logged_body(data) + request_data: Final = types.MappingProxyType( + dict( + ( + *logged_body.items(), + *((("extra_headers", extra_headers),) if extra_headers else ()), + ) + ) + ) + return logged_body, request_data + + class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): def __init__(self) -> None: super().__init__() @@ -1547,6 +1588,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, @@ -1554,11 +1596,10 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): "headers": {"Authorization": f"Bearer {openai_aclient.api_key}"}, "api_base": str(openai_aclient.base_url), "acompletion": True, - "complete_input_dict": data, + "complete_input_dict": logged_body, }, ) - request_data: Final = {**data, "extra_headers": headers} if headers else data raw_response: Final = await openai_aclient.images.with_raw_response.generate( **request_data, timeout=timeout ) @@ -1567,7 +1608,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=stringified_response, ) image_response: Final[ImageResponse] = convert_to_model_response_object( @@ -1582,6 +1623,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj.post_call( input=prompt, api_key=api_key, + additional_args={"complete_input_dict": _image_generation_logged_body(data)}, original_response=str(e), ) raise e @@ -1636,6 +1678,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client=client, ) + logged_body, request_data = _image_generation_request_data(data, headers) + ## LOGGING logging_obj.pre_call( input=prompt, @@ -1644,12 +1688,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): "headers": {"Authorization": f"Bearer {openai_client.api_key}"}, "api_base": openai_client._base_url._uri_reference, # pyright: ignore[reportPrivateUsage] # SDK URL internals "acompletion": True, - "complete_input_dict": data, + "complete_input_dict": logged_body, }, ) ## COMPLETION CALL - request_data: Final = {**data, "extra_headers": headers} if headers else data raw_response: Final = openai_client.images.with_raw_response.generate(**request_data, timeout=timeout) response: Final = raw_response.parse().model_dump() @@ -1657,7 +1700,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, ) image_response: Final[ImageResponse] = convert_to_model_response_object( @@ -1672,7 +1715,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": _image_generation_logged_body(data)}, original_response=str(e), ) raise e @@ -1681,7 +1724,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": _image_generation_logged_body(data)}, original_response=str(e), ) if hasattr(e, "status_code"): diff --git a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index e103412e0c3..5ff18b51cb9 100644 --- a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -6,11 +6,11 @@ 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.common_utils import OpenAIError from litellm.llms.openai.openai import OpenAIChatCompletion from litellm.types.utils import ImageResponse @@ -21,6 +21,21 @@ def _raw_image_response(mock_image_data): raw_response.headers = {"x-request-id": "req-image"} return raw_response +_SECRET_HEADERS = {"Authorization": "Bearer leaked-token", "cf-aig-authorization": "Bearer custom-token"} + + +def _assert_logged_body_omits_credentials(logged_body: object) -> None: + assert isinstance(logged_body, dict) + assert "extra_headers" not in logged_body + assert "Authorization" not in logged_body + assert "cf-aig-authorization" not in logged_body + extra_body = logged_body.get("extra_body") or {} + assert "extra_headers" not in extra_body + assert "Bearer leaked-token" not in str(logged_body) + assert extra_body.get("foo") == "bar" + assert logged_body.get("extra_query") == {"user": "alice"} + assert logged_body.get("size") == "1024x1024" + @pytest.fixture def openai_chat_completions(): @@ -38,9 +53,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 = { @@ -69,9 +82,7 @@ class TestImageGenerationExtraHeaders: _, kwargs = mock_openai_client.images.with_raw_response.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 = { @@ -98,9 +109,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 = { @@ -131,9 +140,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 = { @@ -207,20 +214,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.with_raw_response.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"} @@ -240,6 +241,157 @@ 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"}, + "extra_query": {"user": "alice"}, + } + + if is_async: + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(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=ImageResponse(), + timeout=60.0, + logging_obj=mock_logging_obj, + api_key="test-key", + client=mock_openai_client, + ) + else: + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(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.with_raw_response.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" + assert kwargs.get("extra_query") == {"user": "alice"} + + logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"]["complete_input_dict"] + _assert_logged_body_omits_credentials(logged_body) + + 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.with_raw_response.generate.return_value = _raw_image_response(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.with_raw_response.generate.call_args + assert kwargs.get("extra_headers") == {"cf-aig-authorization": "Bearer from-headers"} + + @pytest.mark.parametrize( + "is_async,raised", + [ + (False, OpenAIError(status_code=400, message="bad image params")), + (False, RuntimeError("forced failure")), + (True, OpenAIError(status_code=400, message="bad image params")), + (True, RuntimeError("forced failure")), + ], + ids=["sync-openai-error", "sync-runtime-error", "async-openai-error", "async-runtime-error"], + ) + @pytest.mark.asyncio + async def test_failure_telemetry_omits_extra_headers( + self, openai_chat_completions, mock_logging_obj, is_async, raised + ): + """post_call on a generate() failure must use the sanitized body, not raw data.""" + mock_openai_client = MagicMock() + mock_openai_client.api_key = "test-key" + mock_openai_client._base_url._uri_reference = "https://api.openai.com" + + optional_params = { + "size": "1024x1024", + "extra_headers": _SECRET_HEADERS, + "extra_body": {"extra_headers": _SECRET_HEADERS, "foo": "bar"}, + "extra_query": {"user": "alice"}, + } + + if is_async: + mock_openai_client.images.with_raw_response.generate = AsyncMock(side_effect=raised) + with pytest.raises(type(raised)): + 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", + headers=_SECRET_HEADERS, + client=mock_openai_client, + ) + else: + mock_openai_client.images.with_raw_response.generate.side_effect = raised + with pytest.raises(OpenAIError): + 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", + headers=_SECRET_HEADERS, + client=mock_openai_client, + ) + + mock_openai_client.images.with_raw_response.generate.assert_called_once() + logged_body = mock_logging_obj.post_call.call_args[1]["additional_args"]["complete_input_dict"] + _assert_logged_body_omits_credentials(logged_body) + class TestImageGenerationEntryPointHeaders: """Test that litellm.image_generation() passes headers through to the OpenAI provider.""" @@ -275,3 +427,4 @@ class TestImageGenerationEntryPointHeaders: mock_openai_client.images.with_raw_response.generate.assert_called_once() _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers + assert "extra_headers" not in (kwargs.get("extra_body") or {})