scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 156

Merge f1c2e169de into 0c98afa780

Browse source
This commit is contained in:
Zsanz3 2026-09-12 14:54:32 -04:00 • committed by GitHub
commit 62cc136cf3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 144 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/llms/openai/openai.py
View file

@ -340,6 +340,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 _embedding_request_without_sdk_defaults(
data: Mapping[str, object], timeout: float | httpx.Timeout
) -> tuple[Mapping[str, object], RequestOptions]:
@ -356,6 +362,37 @@ def _embedding_request_without_sdk_defaults(
return body, options
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_body: Final = types.MappingProxyType(
{key: item for key, item in extra_body_source.items() if key != "extra_headers"}
)
extra_headers: Final = types.MappingProxyType(
{
**_as_mapping(extra_body_source.get("extra_headers")),
**_as_mapping(data.get("extra_headers")),
**_as_mapping(headers),
}
)
extra_query: Final = _as_mapping(data.get("extra_query"))
logged_body: Final = { # 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
}
request_data: Final = types.MappingProxyType(
dict(
(
*logged_body.items(),
*((("extra_headers", extra_headers),) if extra_headers else ()),
*((("extra_query", extra_query),) if extra_query else ()),
*((("extra_body", extra_body),) if extra_body else ()),
)
)
)
return logged_body, request_data
class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
def __init__(self) -> None:
super().__init__()
@ -1415,6 +1452,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,
@ -1422,20 +1460,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
"headers": {"Authorization": f"Bearer {openai_aclient.api_key}"}, # mutable-ok: logged header map
"api_base": str(openai_aclient.base_url),
"acompletion": True,
"complete_input_dict": data,
"complete_input_dict": logged_body,
},
)
request_data: Final = ( # mutable-ok: the OpenAI SDK takes the request body as a dict
{**data, "extra_headers": headers} if headers else data
)
response = await openai_aclient.images.generate(**request_data, timeout=timeout)
stringified_response: Final = response.model_dump()
## LOGGING
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,
)
return convert_to_model_response_object(
@ -1499,6 +1534,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
client=client,
)
logged_body, request_data = _image_generation_request_data(data, headers)
## LOGGING
logging_obj.pre_call(
input=prompt,
@ -1507,14 +1544,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
"headers": {"Authorization": f"Bearer {openai_client.api_key}"},
"api_base": openai_client._base_url._uri_reference,
"acompletion": True,
"complete_input_dict": data,
"complete_input_dict": logged_body,
},
)
## COMPLETION CALL
request_data: Final = ( # mutable-ok: the OpenAI SDK takes the request body as a dict
{**data, "extra_headers": headers} if headers else data
)
_response: Final = openai_client.images.generate(**request_data, timeout=timeout)
response: Final = _response.model_dump()
@ -1522,7 +1556,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,
)
return convert_to_model_response_object(

123
tests/test_litellm/llms/openai/image_generation/test_openai_image_generation_extra_headers.py
View file

@ -6,11 +6,10 @@ 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.openai import OpenAIChatCompletion
@ -30,9 +29,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 = {
@ -61,9 +58,7 @@ class TestImageGenerationExtraHeaders:
_, 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
):
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 = {
@ -90,9 +85,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 = {
@ -121,9 +114,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 = {
@ -193,20 +184,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.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"}
@ -226,6 +211,97 @@ 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"},
}
if is_async:
mock_openai_client.images.generate = AsyncMock(return_value=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=MagicMock(),
timeout=60.0,
logging_obj=mock_logging_obj,
api_key="test-key",
client=mock_openai_client,
)
else:
mock_openai_client.images.generate.return_value = 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.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"
logged_body = mock_logging_obj.pre_call.call_args[1]["additional_args"]["complete_input_dict"]
assert "extra_headers" not in logged_body
assert "extra_headers" not in (logged_body.get("extra_body") or {})
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.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="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.generate.call_args
assert kwargs.get("extra_headers") == {"cf-aig-authorization": "Bearer from-headers"}
class TestImageGenerationEntryPointHeaders:
"""Test that litellm.image_generation() passes headers through to the OpenAI provider."""
@ -259,3 +335,4 @@ class TestImageGenerationEntryPointHeaders:
mock_openai_client.images.generate.assert_called_once()
_, kwargs = mock_openai_client.images.generate.call_args
assert kwargs.get("extra_headers") == test_headers
assert "extra_headers" not in (kwargs.get("extra_body") or {})