mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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:
parent
3bd9c950f8
commit
80b64d53c6
3 changed files with 24 additions and 3 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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="
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue