diff --git a/litellm/images/main.py b/litellm/images/main.py index 7f94835d59b..14de5dd0c4c 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -422,6 +422,7 @@ def image_generation( logging_obj=litellm_logging_obj, timeout=timeout, client=client, + _is_async=aimg_generation is True, ) elif custom_llm_provider == "black_forest_labs": # Route to BFL-specific handler (polling required) diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 7ac440c6f48..811f387982e 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -100,6 +100,32 @@ class BaseImageGenerationConfig(ABC): "ImageVariationConfig implements 'transform_response_image_variation' for image variation models" ) + async def async_transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict[str, object], # mutable-ok: mirrors the sync hook so the default passes the same objects + optional_params: dict[str, object], # mutable-ok: mirrors the sync hook so the default passes the same objects + litellm_params: dict[str, object], # mutable-ok: mirrors the sync hook so the default passes the same objects + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + return self.transform_image_generation_response( # pyright: ignore[reportUnknownMemberType] # sync hook signature is untyped + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + def use_multipart_form_data(self) -> bool: """ Returns True if this provider requires multipart/form-data instead of JSON. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b03b8050299..cb6871cad42 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -44,7 +44,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) -from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, @@ -7464,7 +7464,7 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client - headers: Final = image_generation_provider_config.validate_environment( + headers: Final = await asyncify(image_generation_provider_config.validate_environment)( # pyright: ignore[reportUnknownArgumentType] # sync hook signature is untyped api_key=api_key, headers=image_generation_optional_request_params.get("extra_headers", {}) or {}, model=model, @@ -7529,13 +7529,15 @@ class BaseLLMHTTPHandler: provider_config=image_generation_provider_config, ) - model_response: Final[ImageResponse] = image_generation_provider_config.transform_image_generation_response( + model_response: Final[ + ImageResponse + ] = await image_generation_provider_config.async_transform_image_generation_response( model=model, raw_response=response, model_response=litellm.ImageResponse(), logging_obj=logging_obj, - request_data=data, - optional_params=image_generation_optional_request_params, + request_data=data, # pyright: ignore[reportUnknownArgumentType] # handler params are untyped + optional_params=image_generation_optional_request_params, # pyright: ignore[reportUnknownArgumentType] # handler params are untyped litellm_params=dict(litellm_params), encoding=None, ) diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index ff43e892eab..631646a7fee 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -1,6 +1,8 @@ import asyncio +import itertools import json from datetime import datetime +from collections.abc import Coroutine from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -13,6 +15,7 @@ from openai.types.image import Image import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import CallTypes, StandardLoggingPayload import os @@ -249,122 +252,59 @@ class TestAimlImageGeneration: pytest.fail(f"An exception occurred - {str(e)}") +def _capturing_client(response_json: dict[str, object]) -> tuple[AsyncHTTPHandler, list[httpx.Request]]: + requests: Final[list[httpx.Request]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json=response_json) + + return AsyncHTTPHandler(transport=httpx.MockTransport(handler)), requests + + @pytest.mark.asyncio async def test_aiml_image_generation_with_dynamic_api_key(): - """ - Test that when api_key is passed as a dynamic parameter to aimage_generation, - it gets properly used for AIML provider authentication instead of falling back - to environment variables. + client, requests = _capturing_client( + {"created": 1703658209, "data": [{"url": "https://example.com/generated_image.png"}]} + ) - This test validates the fix for ensuring dynamic API keys are respected - when making image generation requests to the AIML provider. - """ - from unittest.mock import AsyncMock, MagicMock, patch + await litellm.aimage_generation( + prompt="A cute baby sea otter", + model="aiml/flux-pro/v1.1", + api_key="test-dynamic-api-key-12345", + client=client, + ) - import httpx - - # Mock AIML response - mock_aiml_response = { - "created": 1703658209, - "data": [{"url": "https://example.com/generated_image.png"}], - } - - # Track captured arguments - captured_headers = None - captured_url = None - captured_json_data = None - - def capture_post_call(*args, **kwargs): - nonlocal captured_headers, captured_url, captured_json_data - captured_url = kwargs.get("url") or (args[0] if args else None) - captured_headers = kwargs.get("headers", {}) - captured_json_data = kwargs.get("json", {}) - - # Create a mock response - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = mock_aiml_response - mock_response.text = json.dumps(mock_aiml_response) - return mock_response - - # Mock the HTTP client that actually makes the request (sync version for image generation) - with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: - mock_post.side_effect = capture_post_call - - # Test with dynamic api_key - test_api_key = "test-dynamic-api-key-12345" - - response = await litellm.aimage_generation( - prompt="A cute baby sea otter", - model="aiml/flux-pro/v1.1", - api_key=test_api_key, # This should be used instead of env vars - ) - - # Validate the response (mocked response processing might not populate data correctly) - assert response is not None - - # The most important validations: API key and endpoint usage - # These prove that the dynamic API key was properly used - assert captured_headers is not None - assert "Authorization" in captured_headers - assert captured_headers["Authorization"] == f"Bearer {test_api_key}" - print("TESTCAPTURED HEADERS", captured_headers) - # Validate the correct AIML endpoint was called - assert captured_url is not None - assert "api.aimlapi.com" in captured_url - assert "/v1/images/generations" in captured_url - - # Validate the request data - assert captured_json_data is not None - assert captured_json_data["prompt"] == "A cute baby sea otter" - assert captured_json_data["model"] == "flux-pro/v1.1" + assert len(requests) == 1 + assert requests[0].headers["Authorization"] == "Bearer test-dynamic-api-key-12345" + assert requests[0].url.host == "api.aimlapi.com" + assert requests[0].url.path.endswith("/v1/images/generations") + body: Final = json.loads(requests[0].content) + assert body["prompt"] == "A cute baby sea otter" + assert body["model"] == "flux-pro/v1.1" @pytest.mark.asyncio async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape(): - """End-to-end check that ``aiml/openai/gpt-image-2`` keeps the upstream - OpenAI request shape (``size``/``n``/``response_format``) instead of - being remapped to the AI/ML flux schema (``image_size``/``num_images``/ - ``output_format``), and hits the correct upstream model name. - """ - import json as _json - from unittest.mock import MagicMock, patch + client, requests = _capturing_client( + {"created": 1703658209, "data": [{"url": "https://example.com/gpt-image-2.png"}]} + ) - mock_aiml_response = { - "created": 1703658209, - "data": [{"url": "https://example.com/gpt-image-2.png"}], - } + await litellm.aimage_generation( + prompt="A T-Rex relaxing on a beach", + model="aiml/openai/gpt-image-2", + api_key="test-key-mocked-no-credits-needed", + size="1024x1536", + quality="high", + response_format="b64_json", + n=1, + client=client, + ) - captured = {} - - def capture_post_call(*args, **kwargs): - captured["url"] = kwargs.get("url") or (args[0] if args else None) - captured["headers"] = kwargs.get("headers", {}) - captured["json"] = kwargs.get("json", {}) - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.json.return_value = mock_aiml_response - mock_response.text = _json.dumps(mock_aiml_response) - return mock_response - - with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: - mock_post.side_effect = capture_post_call - - await litellm.aimage_generation( - prompt="A T-Rex relaxing on a beach", - model="aiml/openai/gpt-image-2", - api_key="test-key-mocked-no-credits-needed", - size="1024x1536", - quality="high", - response_format="b64_json", - n=1, - ) - - assert captured["url"] is not None - assert "api.aimlapi.com" in captured["url"] - assert "/v1/images/generations" in captured["url"] - - body = captured["json"] + assert len(requests) == 1 + assert requests[0].url.host == "api.aimlapi.com" + assert requests[0].url.path.endswith("/v1/images/generations") + body: Final = json.loads(requests[0].content) assert body["model"] == "openai/gpt-image-2" assert body["prompt"] == "A T-Rex relaxing on a beach" assert body["size"] == "1024x1536" @@ -405,3 +345,33 @@ 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 +async def test_aimage_generation_runs_vertex_requests_concurrently_on_the_event_loop() -> None: + response_json: Final = { + "candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}], + } + arrivals: Final = itertools.count(1) + both_in_flight: Final = asyncio.Event() + + async def handler(request: httpx.Request) -> httpx.Response: + if next(arrivals) == 2: + both_in_flight.set() + await both_in_flight.wait() + return httpx.Response(status_code=200, json=response_json) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(handler)) + + def generate() -> Coroutine[object, object, litellm.ImageResponse]: + return litellm.aimage_generation( # pyright: ignore[reportUnknownMemberType] # aimage_generation kwargs are untyped + model="vertex_ai/gemini-2.5-flash-image", + prompt="a red circle", + api_base="http://localhost:1/generateContent", + vertex_location="us-central1", + client=client, + ) + + results: Final = await asyncio.wait_for(asyncio.gather(generate(), generate()), timeout=5) + + assert [result.data[0].b64_json if result.data else None for result in results] == ["aW1n", "aW1n"] diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 23c142b61ab..c88933eeb99 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -64,6 +64,7 @@ from litellm.responses.streaming_iterator import ( ) from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @@ -4902,6 +4903,62 @@ async def test_async_image_generation_handler_records_upstream_response_headers( _assert_upstream_headers_recorded(response) +class _OffLoopImageGenerationConfig(_HeaderImageGenerationConfig): + def validate_environment( + self, + headers: dict[str, object], + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + litellm_params: dict[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: + try: + asyncio.get_running_loop() + except RuntimeError: + return {"Authorization": "Bearer from-worker-thread"} + raise AssertionError("validate_environment ran on the event loop") + + def transform_image_generation_response(self, *args: object, **kwargs: object) -> ImageResponse: + raise AssertionError("sync transform_image_generation_response ran on the event loop") + + async def async_transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: object, + request_data: dict[str, object], + optional_params: dict[str, object], + litellm_params: dict[str, object], + encoding: object, + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["b64_json"])]) + + +@pytest.mark.asyncio +async def test_async_image_generation_handler_keeps_provider_auth_and_polling_off_the_event_loop(): + seen_authorization: Final[list[str | None]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen_authorization.append(request.headers.get("Authorization")) + return httpx.Response(200, json={"b64_json": "abc"}) + + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + response: Final = await BaseLLMHTTPHandler().async_image_generation_handler( # pyright: ignore[reportUnknownMemberType] # handler params are untyped + client=client, + **{**_image_generation_call_kwargs(), "image_generation_provider_config": _OffLoopImageGenerationConfig()}, # pyright: ignore[reportUnknownArgumentType] # call kwargs helper is untyped + ) + + assert response.data[0].b64_json == "abc" # pyright: ignore[reportUnknownMemberType] # ImageObject fields are partially untyped + assert seen_authorization == ["Bearer from-worker-thread"] + + class _HeaderTextToSpeechConfig(BaseTextToSpeechConfig): def get_supported_openai_params(self, model): return [] diff --git a/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation.py b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation.py index 00c0a75890d..9621e0dd67f 100644 --- a/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation.py +++ b/tests/unit/llms/fal_ai/image_generation/test_fal_ai_image_generation.py @@ -1,8 +1,11 @@ -from unittest.mock import MagicMock, patch +import json +from typing import Final +import httpx import pytest from litellm import aimage_generation +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -22,78 +25,40 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco ], ) @pytest.mark.asyncio -async def test_fal_ai_image_generation_basic(model, expected_endpoint): - """ - Test that fal_ai image generation constructs correct request body and URL. +async def test_fal_ai_image_generation_basic(model: str, expected_endpoint: str) -> None: + requests: Final[list[httpx.Request]] = [] - Validates: - - Correct API endpoint URL construction - - Proper request body format with prompt - - Correct Authorization header format - """ - captured_url = None - captured_json_data = None - captured_headers = None - - def capture_post_call(*args, **kwargs): - nonlocal captured_url, captured_json_data, captured_headers - - captured_url = args[0] if args else kwargs.get("url") - captured_json_data = kwargs.get("json") - captured_headers = kwargs.get("headers") - - # Mock response with fal.ai format - mock_response = MagicMock() - mock_response.status_code = 200 - mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = { - "images": [ - { - "url": "https://example.com/generated-image.png", - "width": 1024, - "height": 768, - "content_type": "image/jpeg", - } - ], - "seed": 42, - } - - return mock_response - - with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: - mock_post.side_effect = capture_post_call - - test_api_key = "test-fal-ai-key-12345" - test_prompt = "A cute baby sea otter" - - response = await aimage_generation( - model=model, - prompt=test_prompt, - api_key=test_api_key, + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + 200, + json={ + "images": [ + { + "url": "https://example.com/generated-image.png", + "width": 1024, + "height": 768, + "content_type": "image/jpeg", + } + ], + "seed": 42, + }, ) - # Validate response - assert response is not None - assert hasattr(response, "data") - assert response.data is not None - assert len(response.data) > 0 + response: Final = await aimage_generation( + model=model, + prompt="A cute baby sea otter", + api_key="test-fal-ai-key-12345", + client=AsyncHTTPHandler(transport=httpx.MockTransport(handler)), + ) - # Validate URL - assert captured_url is not None - assert "fal.run" in captured_url - assert expected_endpoint in captured_url - print(f"Validated URL: {captured_url}") - - # Validate headers - assert captured_headers is not None - assert "Authorization" in captured_headers - assert captured_headers["Authorization"] == f"Key {test_api_key}" - print(f"Validated headers: {captured_headers}") - - # Validate request body - assert captured_json_data is not None - assert captured_json_data["prompt"] == test_prompt - print(f"Validated request body: {captured_json_data}") + assert response.data + assert response.data[0].url == "https://example.com/generated-image.png" + assert len(requests) == 1 + assert requests[0].url.host == "fal.run" + assert expected_endpoint in requests[0].url.path + assert requests[0].headers["Authorization"] == "Key test-fal-ai-key-12345" + assert json.loads(requests[0].content)["prompt"] == "A cute baby sea otter" @pytest.fixture(autouse=True)