diff --git a/litellm/utils.py b/litellm/utils.py index 54de4ebbbb2..95acebc4c0a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4707,7 +4707,7 @@ def get_optional_params_image_gen( custom_handler: Final = ( ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) - if custom_llm_provider not in LlmProvidersSet + if custom_llm_provider not in ("bedrock", "vertex_ai", "azure_ai", "black_forest_labs", "litellm_proxy") else None ) if provider_config is not None: diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 79678e3fe43..86f426834d0 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -792,3 +792,22 @@ async def test_custom_images_map_handler_errors_after_forwarding_params( with pytest.raises(litellm.RateLimitError, match="image quota exhausted"): call(**request, **({"image": b"input-image"} if operation == "edit" else {})) assert custom_image_params_handler.requests[0]["user"] == "reject" + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("with_params", [False, True]) +@pytest.mark.asyncio +async def test_custom_images_receive_params_for_builtin_names_without_image_support( + custom_image_params_handler: CustomImageParamsLLM, monkeypatch: pytest.MonkeyPatch, + is_async: bool, with_params: bool, +) -> None: + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "anthropic", "custom_handler": custom_image_params_handler}] + ) + litellm.utils.custom_llm_setup() + supplied: Final = {"n": 2, "size": "1024x1024"} if with_params else {} + request: Final = dict(model="anthropic/custom-image-model", prompt="draw", **supplied) + result: Final = await litellm.aimage_generation(**request) if is_async else litellm.image_generation(**request) + received: Final = custom_image_params_handler.requests[0] + assert {key: received[key] for key in ("n", "size") if key in received} == supplied + assert result.data[0].b64_json == "aW1hZ2U=" diff --git a/tests/unit/test_utils_get_optional_params.py b/tests/unit/test_utils_get_optional_params.py index dd16f2f3ef6..7edffc7972e 100644 --- a/tests/unit/test_utils_get_optional_params.py +++ b/tests/unit/test_utils_get_optional_params.py @@ -542,7 +542,9 @@ def test_get_optional_params_image_gen(): @pytest.mark.parametrize("with_params", [False, True]) -def test_image_params_keep_builtin_mapping_when_custom_name_collides(monkeypatch, with_params): +def test_image_params_keep_builtin_mapping_when_custom_name_collides( + monkeypatch: pytest.MonkeyPatch, with_params: bool +) -> None: monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "aiml", "custom_handler": CustomLLM()}]) supplied: Final = {"n": 2, "size": "1024x1024"} if with_params else {} result: Final = get_optional_params_image_gen( @@ -551,7 +553,7 @@ def test_image_params_keep_builtin_mapping_when_custom_name_collides(monkeypatch assert result == ({"num_images": 2, "image_size": {"width": 1024, "height": 1024}} if with_params else {}) -def test_image_params_keep_builtin_validation_when_custom_name_collides(monkeypatch): +def test_image_params_keep_builtin_validation_when_custom_name_collides(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "aiml", "custom_handler": CustomLLM()}]) with pytest.raises(litellm.UnsupportedParamsError, match="quality"): get_optional_params_image_gen(