From 92d8b16eb848ffcf5d7a6989e8f0037377c487c4 Mon Sep 17 00:00:00 2001 From: yuneng Date: Fri, 9 Oct 2026 21:59:17 +0000 Subject: [PATCH 1/5] fix(images): run llm_http_handler image generation on the event loop --- litellm/images/main.py | 1 + tests/unit/images/test_main.py | 99 +++++++++++++++++++++++++++++++--- 2 files changed, 93 insertions(+), 7 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index af01e051211..ff822a522f8 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -403,6 +403,7 @@ def image_generation( logging_obj=litellm_logging_obj, timeout=timeout, client=client, + _is_async=aimg_generation, ) elif custom_llm_provider == "black_forest_labs": # Route to BFL-specific handler (polling required) diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 189f996c0cc..bfdac96da74 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -1,7 +1,10 @@ import asyncio +import concurrent.futures +import itertools import json +import threading from datetime import datetime -from typing import Final +from typing import Coroutine, Final, Protocol, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -12,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 @@ -212,7 +216,7 @@ async def test_aiml_image_generation_with_dynamic_api_key(): captured_url = None captured_json_data = None - def capture_post_call(*args, **kwargs): + async 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", {}) @@ -225,8 +229,10 @@ async def test_aiml_image_generation_with_dynamic_api_key(): 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: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: mock_post.side_effect = capture_post_call # Test with dynamic api_key @@ -266,7 +272,7 @@ async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape(): ``output_format``), and hits the correct upstream model name. """ import json as _json - from unittest.mock import MagicMock, patch + from unittest.mock import AsyncMock, MagicMock, patch mock_aiml_response = { "created": 1703658209, @@ -275,7 +281,7 @@ async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape(): captured = {} - def capture_post_call(*args, **kwargs): + async 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", {}) @@ -285,7 +291,10 @@ async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape(): mock_response.text = _json.dumps(mock_aiml_response) return mock_response - with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: mock_post.side_effect = capture_post_call await litellm.aimage_generation( @@ -343,3 +352,79 @@ 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 + + +def test_aimage_generation_runs_vertex_requests_on_the_event_loop_not_the_executor( + respx_mock: respx.MockRouter, +) -> None: + api_base: Final = "http://localhost:12348/generateContent" + response_json: Final = { + "candidates": [{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": "aW1n"}}]}}], + } + arrivals: Final = itertools.count(1) + released: Final = threading.Event() + + class AImageGeneration(Protocol): + def __call__( + self, + *, + model: str, + prompt: str, + api_base: str, + vertex_location: str, + client: AsyncHTTPHandler, + ) -> Coroutine[object, object, litellm.ImageResponse]: ... + + aimage_generation: Final[AImageGeneration] = cast(AImageGeneration, getattr(litellm, "aimage_generation")) + + def held_sync(request: httpx.Request) -> httpx.Response: + if next(arrivals) == 2: + released.set() + released.wait() + return httpx.Response(status_code=200, json=response_json) + + respx_mock.post(api_base).mock(side_effect=held_sync) + + async def scenario() -> tuple[litellm.ImageResponse, litellm.ImageResponse]: + asyncio.get_running_loop().set_default_executor(concurrent.futures.ThreadPoolExecutor(max_workers=1)) + in_flight: Final = asyncio.Event() + + async def handler(request: httpx.Request) -> httpx.Response: + if next(arrivals) == 2: + in_flight.set() + await in_flight.wait() + return httpx.Response(status_code=200, json=response_json) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(handler)) + try: + return await asyncio.wait_for( + asyncio.gather( + aimage_generation( + model="vertex_ai/gemini-2.5-flash-image", + prompt="a red circle", + api_base=api_base, + vertex_location="us-central1", + client=client, + ), + aimage_generation( + model="vertex_ai/gemini-2.5-flash-image", + prompt="a red circle", + api_base=api_base, + vertex_location="us-central1", + client=client, + ), + ), + timeout=5, + ) + finally: + released.set() + + results: Final = asyncio.run(scenario()) + first: Final = results[0] + second: Final = results[1] + assert isinstance(first, litellm.ImageResponse) + assert first.data is not None + assert cast(str | None, getattr(first.data[0], "b64_json")) == "aW1n" + assert isinstance(second, litellm.ImageResponse) + assert second.data is not None + assert cast(str | None, getattr(second.data[0], "b64_json")) == "aW1n" From 2e472a9cd0e718717be891f7b2b5559eb19105cf Mon Sep 17 00:00:00 2001 From: yuneng Date: Fri, 9 Oct 2026 22:05:45 +0000 Subject: [PATCH 2/5] test(images): patch the async http handler in the fal image generation test --- tests/unit/images/test_main.py | 22 ++++++------------- .../test_fal_ai_image_generation.py | 9 +++++--- 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index bfdac96da74..6f64cb11138 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -4,7 +4,7 @@ import itertools import json import threading from datetime import datetime -from typing import Coroutine, Final, Protocol, cast +from typing import Callable, Coroutine, Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -364,18 +364,10 @@ def test_aimage_generation_runs_vertex_requests_on_the_event_loop_not_the_execut arrivals: Final = itertools.count(1) released: Final = threading.Event() - class AImageGeneration(Protocol): - def __call__( - self, - *, - model: str, - prompt: str, - api_base: str, - vertex_location: str, - client: AsyncHTTPHandler, - ) -> Coroutine[object, object, litellm.ImageResponse]: ... - - aimage_generation: Final[AImageGeneration] = cast(AImageGeneration, getattr(litellm, "aimage_generation")) + aimage_generation: Final = cast( + Callable[..., Coroutine[object, object, litellm.ImageResponse]], + getattr(litellm, "aimage_generation"), + ) def held_sync(request: httpx.Request) -> httpx.Response: if next(arrivals) == 2: @@ -424,7 +416,7 @@ def test_aimage_generation_runs_vertex_requests_on_the_event_loop_not_the_execut second: Final = results[1] assert isinstance(first, litellm.ImageResponse) assert first.data is not None - assert cast(str | None, getattr(first.data[0], "b64_json")) == "aW1n" + assert first.data[0].b64_json == "aW1n" assert isinstance(second, litellm.ImageResponse) assert second.data is not None - assert cast(str | None, getattr(second.data[0], "b64_json")) == "aW1n" + assert second.data[0].b64_json == "aW1n" 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..9b285959d8f 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,4 +1,4 @@ -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -35,7 +35,7 @@ async def test_fal_ai_image_generation_basic(model, expected_endpoint): captured_json_data = None captured_headers = None - def capture_post_call(*args, **kwargs): + async def capture_post_call(*args, **kwargs): nonlocal captured_url, captured_json_data, captured_headers captured_url = args[0] if args else kwargs.get("url") @@ -60,7 +60,10 @@ async def test_fal_ai_image_generation_basic(model, expected_endpoint): return mock_response - with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: mock_post.side_effect = capture_post_call test_api_key = "test-fal-ai-key-12345" From adb8b0a3e0f28c69161a5c254de548e016ce723a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 9 Oct 2026 15:22:50 -0700 Subject: [PATCH 3/5] refactor(images): pass a typed bool as the image handler async flag --- litellm/images/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index ff822a522f8..c94a62047fb 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -403,7 +403,7 @@ def image_generation( logging_obj=litellm_logging_obj, timeout=timeout, client=client, - _is_async=aimg_generation, + _is_async=aimg_generation is True, ) elif custom_llm_provider == "black_forest_labs": # Route to BFL-specific handler (polling required) From 113d98fbd3e52eb1d37e36e7bcb23c66da193490 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 9 Oct 2026 15:35:39 -0700 Subject: [PATCH 4/5] fix(images): keep provider auth and response polling off the event loop for async image generation --- .../image_generation/transformation.py | 26 +++++++ litellm/llms/custom_httpx/llm_http_handler.py | 12 +-- tests/unit/images/test_main.py | 77 +++++-------------- .../custom_httpx/test_llm_http_handler.py | 57 ++++++++++++++ 4 files changed, 109 insertions(+), 63 deletions(-) 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 4a9b8fee2e8..ef495c4de76 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -42,7 +42,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, @@ -7292,7 +7292,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, @@ -7357,13 +7357,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 6f64cb11138..4090ae208e0 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -1,10 +1,9 @@ import asyncio -import concurrent.futures import itertools import json -import threading from datetime import datetime -from typing import Callable, Coroutine, Final, cast +from collections.abc import Coroutine +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -354,69 +353,31 @@ async def test_azure_image_generation_request_body(): assert request_json == expected_body -def test_aimage_generation_runs_vertex_requests_on_the_event_loop_not_the_executor( - respx_mock: respx.MockRouter, -) -> None: - api_base: Final = "http://localhost:12348/generateContent" +@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) - released: Final = threading.Event() + both_in_flight: Final = asyncio.Event() - aimage_generation: Final = cast( - Callable[..., Coroutine[object, object, litellm.ImageResponse]], - getattr(litellm, "aimage_generation"), - ) - - def held_sync(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx.Request) -> httpx.Response: if next(arrivals) == 2: - released.set() - released.wait() + both_in_flight.set() + await both_in_flight.wait() return httpx.Response(status_code=200, json=response_json) - respx_mock.post(api_base).mock(side_effect=held_sync) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(handler)) - async def scenario() -> tuple[litellm.ImageResponse, litellm.ImageResponse]: - asyncio.get_running_loop().set_default_executor(concurrent.futures.ThreadPoolExecutor(max_workers=1)) - in_flight: Final = asyncio.Event() + 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, + ) - async def handler(request: httpx.Request) -> httpx.Response: - if next(arrivals) == 2: - in_flight.set() - await in_flight.wait() - return httpx.Response(status_code=200, json=response_json) + results: Final = await asyncio.wait_for(asyncio.gather(generate(), generate()), timeout=5) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(handler)) - try: - return await asyncio.wait_for( - asyncio.gather( - aimage_generation( - model="vertex_ai/gemini-2.5-flash-image", - prompt="a red circle", - api_base=api_base, - vertex_location="us-central1", - client=client, - ), - aimage_generation( - model="vertex_ai/gemini-2.5-flash-image", - prompt="a red circle", - api_base=api_base, - vertex_location="us-central1", - client=client, - ), - ), - timeout=5, - ) - finally: - released.set() - - results: Final = asyncio.run(scenario()) - first: Final = results[0] - second: Final = results[1] - assert isinstance(first, litellm.ImageResponse) - assert first.data is not None - assert first.data[0].b64_json == "aW1n" - assert isinstance(second, litellm.ImageResponse) - assert second.data is not None - assert second.data[0].b64_json == "aW1n" + assert [image.b64_json for result in results for image in result.data or []] == ["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 25d2d11b5ca..8e5bc756ffa 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -59,6 +59,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 @@ -4645,6 +4646,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 [] From 8e5ce6d66c9cc677c9b7c2e50c3caa0637b6cc34 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 9 Oct 2026 16:36:34 -0700 Subject: [PATCH 5/5] test(images): inject mock-transport clients instead of patching the async handler --- tests/unit/images/test_main.py | 156 +++++------------- .../test_fal_ai_image_generation.py | 106 ++++-------- 2 files changed, 78 insertions(+), 184 deletions(-) diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 4090ae208e0..a9353c12253 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -190,127 +190,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 - - async 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 - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) 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 AsyncMock, 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 = {} - - async 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.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) 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" @@ -380,4 +312,4 @@ async def test_aimage_generation_runs_vertex_requests_concurrently_on_the_event_ results: Final = await asyncio.wait_for(asyncio.gather(generate(), generate()), timeout=5) - assert [image.b64_json for result in results for image in result.data or []] == ["aW1n", "aW1n"] + assert [result.data[0].b64_json if result.data else None for result in results] == ["aW1n", "aW1n"] 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 9b285959d8f..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 AsyncMock, 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,81 +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 - - async 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.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) 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)