diff --git a/litellm/images/main.py b/litellm/images/main.py index 29bac3ffe05..9c88e737a40 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import importlib -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import TYPE_CHECKING, Final, Literal, cast, overload @@ -85,6 +85,47 @@ def _get_ImageEditRequestUtils() -> type["ImageEditRequestUtils"]: return _ImageEditRequestUtils_cache +def _chatgpt_image_params(kwargs: Mapping[str, object]) -> GenericLiteLLMParams: + from litellm.responses.utils import ResponsesAPIRequestUtils + + headers: Final = ResponsesAPIRequestUtils.merge_client_forwarded_headers( + TypeAdapter[dict[str, object] | None](dict[str, object] | None).validate_python(kwargs.get("extra_headers")), + TypeAdapter[dict[str, str] | None](dict[str, str] | None).validate_python(kwargs.get("headers")), + ) + return GenericLiteLLMParams.model_validate({**kwargs, "extra_headers": headers}) + + +def _chatgpt_image_headers(model: str, params: GenericLiteLLMParams, headers: Mapping[str, str]) -> Mapping[str, str]: + from litellm.llms.chatgpt.common_utils import get_chatgpt_request_credential + from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig + + credential: Final = get_chatgpt_request_credential(params) + credential_headers: Final = ( + { + "Authorization": f"Bearer {credential.access_token}", + **({"ChatGPT-Account-Id": credential.account_id} if credential.account_id else {}), + } + if credential is not None + else {} + ) + explicit_account_id: Final = next( + (value for name, value in headers.items() if name.lower() == "chatgpt-account-id"), None + ) + auth_headers: Final = ChatGPTResponsesAPIConfig().validate_environment( + headers=credential_headers, model=model, litellm_params=params + ) + return { + **{k: v for k, v in auth_headers.items() if k.lower() not in ("content-type", "accept")}, + **{ + k: v + for k, v in headers.items() + if k.lower() not in ("authorization", "chatgpt-account-id", "cookie", "x-litellm-api-key") + }, + **({"ChatGPT-Account-Id": explicit_account_id} if credential is None and explicit_account_id else {}), + "Accept": "application/json", + } + + ##### Image Generation ####################### @client async def aimage_generation(*args, **kwargs) -> ImageResponse: @@ -121,6 +162,12 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: model=model, custom_llm_provider=provider_hints.get("custom_llm_provider"), api_base=provider_hints.get("api_base"), + litellm_params=( + _chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python(kwargs)) + if provider_hints.get("custom_llm_provider") == "chatgpt" + or (isinstance(model, str) and model.startswith("chatgpt/")) + else None + ), ) # Await normally @@ -240,12 +287,18 @@ def image_generation( if extra_headers is not None: headers.update(extra_headers) model_response: ImageResponse = litellm.utils.ImageResponse() + chatgpt_params: Final = ( + _chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python({**kwargs, "api_key": api_key})) + if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/")) + else None + ) dynamic_api_key: str | None = None if model is not None or custom_llm_provider is not None: model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, + litellm_params=chatgpt_params, ) else: model = "dall-e-2" @@ -268,7 +321,10 @@ def image_generation( "size", "style", ] - non_default_params: Final = filter_out_litellm_params(kwargs, excluding=openai_params) + non_default_params: Final = filter_out_litellm_params( + kwargs, + excluding=[*openai_params, "extra_headers", "headers"] if chatgpt_params is not None else openai_params, + ) image_generation_config: BaseImageGenerationConfig | None = None if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): @@ -490,7 +546,7 @@ def image_generation( or custom_llm_provider == LlmProviders.LITELLM_PROXY.value or custom_llm_provider in litellm.openai_compatible_providers ): - if extra_headers is not None: + if extra_headers is not None and chatgpt_params is None: optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) organization: Final[str | None] = kwargs.get("organization", None) @@ -506,7 +562,17 @@ def image_generation( organization=organization, aimg_generation=aimg_generation, client=client, - headers=headers, + headers=( + dict( + _chatgpt_image_headers( + model, + chatgpt_params.model_copy(update={"api_key": dynamic_api_key or api_key}), + TypeAdapter(dict[str, str]).validate_python(headers), + ) + ) + if chatgpt_params is not None + else headers + ), litellm_params=litellm_params_dict, ) elif custom_llm_provider == "bedrock": diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index ff43e892eab..38af99ff505 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -1,6 +1,8 @@ import asyncio +import base64 import json from datetime import datetime +from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -405,3 +407,56 @@ async def test_azure_image_generation_request_body(): call_args = mock_post.call_args request_json = call_args.kwargs.get("json", {}) assert request_json == expected_body + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("transport", ["headers", "extra_headers", "scoped", "api_key"]) +@pytest.mark.parametrize("saved_key", [None, "fake-saved-gateway-key"]) +async def test_chatgpt_images_forward_client_login_and_workspace( + asynchronous: bool, + transport: str, + saved_key: str | None, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + claims: Final = {"https://api.openai.com/auth": {"chatgpt_account_id": "client-account"}} + payload: Final = base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + token: Final = f"test.{payload}.test" + headers: Final = {"aUtHoRiZaTiOn": f"Bearer {token}", "ChatGPT-Account-Id": "selected-workspace"} + credentials: Final = ( + {"provider_specific_header": {"custom_llm_provider": "chatgpt", "extra_headers": headers}} + if transport == "scoped" + else {"api_key": token, "extra_headers": {"cHaTgPt-AcCoUnT-Id": "selected-workspace"}} + if transport == "api_key" + else {transport: {**headers, "cookie": "fake-cookie", "x-litellm-api-key": "fake-gateway-key"}} + ) + route: Final = respx_mock.post("https://chatgpt.com/backend-api/codex/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]}) + ) + params: Final = { + "model": "chatgpt/gpt-image-2", + "prompt": "A blue circle", + "api_key": saved_key, + **( + { + "client": AsyncOpenAI( + api_key="fake", base_url="https://chatgpt.com/backend-api/codex", http_client=httpx.AsyncClient() + ) + } + if asynchronous + else {} + ), + **credentials, + } + response: Final = await litellm.aimage_generation(**params) if asynchronous else litellm.image_generation(**params) + assert response.data[0].b64_json == "aW1n" + request: Final = route.calls.last.request + assert request.headers["authorization"] == f"Bearer {token}" + assert request.headers["chatgpt-account-id"] == "selected-workspace" + assert request.headers.get_list("content-type") == ["application/json"] + assert "cookie" not in request.headers + assert "x-litellm-api-key" not in request.headers + assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"}