fix: preserve custom image params for non-image built-in names

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-11-06 19:09:33 +00:00
parent 3bd9c950f8
commit 80b64d53c6
3 changed files with 24 additions and 3 deletions

View file

@ -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:

View file

@ -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="

View file

@ -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(