scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(image_edit): await an async transform hook and hide the URL policy verdict from callers

Browse source 
The image edit handler now awaits BaseImageEditConfig.async_transform_image_edit_request, and
Black Forest Labs overrides it so URL images and masks download through async_safe_get instead of
the blocking safe_get on the event loop. Rejected image fetches raise a fixed policy message with
the user_url_allowed_hosts hint rather than echoing the resolver's verdict (resolved IP, DNS
failure) back to the caller. The test fixture also fails any request-path call of the sync
convert_url_to_base64 so a regression cannot pass unnoticed.
This commit is contained in:
mateo-berri 2026-09-04 23:20:21 -07:00
parent 5c8016e99e
commit 0ee3bec046
8 changed files with 379 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/litellm_core_utils/prompt_templates/image_handling.py
View file

@ -78,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str:
return result
def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict)
return litellm.ImageFetchError(
"Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; "
f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}"
)
async def async_convert_url_to_base64(url: str) -> str:
if url.startswith("data:") and ";base64," in url:
return url
@ -100,7 +108,7 @@ async def async_convert_url_to_base64(url: str) -> str:
except litellm.ImageFetchError:
raise
except SSRFError as e:
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e
raise _url_policy_rejection(url, e) from e
except Exception:
pass
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
@ -128,7 +136,7 @@ def convert_url_to_base64(url: str) -> str:
except litellm.ImageFetchError:
raise
except SSRFError as e:
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL. {e} url={url}") from e
raise _url_policy_rejection(url, e) from e
except Exception as e:
verbose_logger.exception(e)
raise litellm.ImageFetchError(

19
litellm/llms/base_llm/image_edit/transformation.py
View file

@ -1,5 +1,6 @@
import types
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any
import httpx
@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC):
) -> tuple[dict, RequestFiles]:
pass
async def async_transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict, RequestFiles]:
return self.transform_image_edit_request(
model=model,
prompt=prompt,
image=image,
image_edit_optional_request_params=dict(image_edit_optional_request_params),
litellm_params=litellm_params,
headers=dict(headers),
)
def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict:
"""
Last pass on the request dict after ``transform_image_edit_request``, using the

102
litellm/llms/black_forest_labs/image_edit/transformation.py
View file

@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/
import base64
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -16,7 +17,7 @@ from httpx._types import RequestFiles
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
@ -37,6 +38,22 @@ else:
LiteLLMLoggingObj = Any
_BFL_REQUEST_PARAMS: Final = (
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
)
class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"""
Configuration for Black Forest Labs image editing.
@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
BFL-specific params are passed through directly.
"""
optional_params: Final[dict[str, Any]] = {}
# Pass through BFL-specific params
bfl_params: Final = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
# Kontext-specific
"aspect_ratio",
# Fill/Inpaint-specific
"steps",
"guidance",
"grow_mask",
# Expand-specific
"top",
"bottom",
"left",
"right",
]
# Convert TypedDict to regular dict for access
params_dict: Final = dict(image_edit_optional_params)
for param in bfl_params:
if param in params_dict:
value = params_dict[param]
if value is not None:
optional_params[param] = value
params: Final[Mapping[str, object]] = image_edit_optional_params
for param in _BFL_REQUEST_PARAMS:
if (value := params.get(param)) is not None:
optional_params[param] = value
# Set default output format
if "output_format" not in optional_params:
@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"input_image": b64_image,
}
# Add optional params (only BFL-recognized parameters)
bfl_request_params: Final = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
]
for key, value in image_edit_optional_request_params.items():
if key in bfl_request_params and value is not None:
if key in _BFL_REQUEST_PARAMS and value is not None:
request_body[key] = value
# Handle mask if provided (for inpainting)
@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
# BFL uses JSON, not multipart - return empty files
return request_body, []
return request_body, ()
async def async_transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict, RequestFiles]:
downloaded_image: Final = await self._fetch_remote_image(image)
downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask"))
return self.transform_image_edit_request(
model=model,
prompt=prompt,
image=image if downloaded_image is None else downloaded_image,
image_edit_optional_request_params=(
dict(image_edit_optional_request_params)
if downloaded_mask is None
else {**image_edit_optional_request_params, "mask": downloaded_mask}
),
litellm_params=litellm_params,
headers=dict(headers),
)
async def _fetch_remote_image(self, image: object) -> bytes | None:
candidate: Final = image[0] if isinstance(image, list) and image else image
if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")):
return None
response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0)
response.raise_for_status()
return response.content
def transform_image_edit_response(
self,

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -6787,7 +6787,7 @@ class BaseLLMHTTPHandler:
litellm_params=dict(litellm_params),
)
data, files = image_edit_provider_config.transform_image_edit_request(
data, files = await image_edit_provider_config.async_transform_image_edit_request(
model=model,
image=image,
prompt=prompt,

10
tests/test_litellm/conftest.py
View file

@ -607,7 +607,8 @@ ONE_PIXEL_PNG = base64.b64decode(
@pytest.fixture
def async_only_image_fetch(monkeypatch):
from litellm.litellm_core_utils.prompt_templates import image_handling
from litellm.litellm_core_utils.prompt_templates import factory, image_handling
from litellm.llms.gemini.chat import transformation as gemini_chat_transformation
fetch = SimpleNamespace(
fetched=[],
@ -627,6 +628,13 @@ def async_only_image_fetch(monkeypatch):
request=httpx.Request("GET", url),
)
def forbid_sync_convert(url, *args, **kwargs):
if url.startswith(("http://", "https://")):
raise litellm.ImageFetchError(f"sync convert_url_to_base64 ran on the request path: {url}")
return url
monkeypatch.setattr(image_handling, "safe_get", forbid_sync_fetch)
monkeypatch.setattr(image_handling, "async_safe_get", serve_png)
for module in (image_handling, factory, gemini_chat_transformation):
monkeypatch.setattr(module, "convert_url_to_base64", forbid_sync_convert)
return fetch

65
tests/test_litellm/litellm_core_utils/test_image_handling.py
View file

@ -428,36 +428,61 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail
assert time.perf_counter() - started < 1
async def test_async_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch):
_SSRF_VERDICTS = (
"URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, "
"add the host to `user_url_allowed_hosts` in general_settings.",
"DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known",
"No addresses found for 'internal.example'",
)
def _assert_one_verdict_free_message(messages, url):
assert len(set(messages)) == 1
assert "10.0.0.8" not in messages[0]
assert "DNS" not in messages[0]
assert "No addresses" not in messages[0]
assert "user_url_allowed_hosts" in messages[0]
assert url in messages[0]
async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
attempts = []
async def block(client, url, **kwargs):
attempts.append(url)
raise SSRFError("URL targets a blocked address (10.0.0.8)")
monkeypatch.setattr(image_handling, "async_safe_get", block)
messages = []
url = f"http://internal.example/{uuid.uuid4()}.png"
with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"):
await async_convert_url_to_base64(url)
for verdict in _SSRF_VERDICTS:
assert attempts == [url]
async def block(client, fetched_url, verdict=verdict, **kwargs):
attempts.append(fetched_url)
raise SSRFError(verdict)
monkeypatch.setattr(image_handling, "async_safe_get", block)
with pytest.raises(litellm.ImageFetchError) as raised:
await async_convert_url_to_base64(url)
messages.append(raised.value.message)
assert attempts == [url] * len(_SSRF_VERDICTS)
_assert_one_verdict_free_message(messages, url)
def test_convert_url_to_base64_reports_a_blocked_url_without_retrying(monkeypatch):
def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch):
attempts = []
def block(client, url, **kwargs):
attempts.append(url)
raise SSRFError("URL targets a blocked address (10.0.0.8)")
monkeypatch.setattr(image_handling, "safe_get", block)
messages = []
url = f"http://internal.example/{uuid.uuid4()}.png"
with pytest.raises(litellm.ImageFetchError, match=r"blocked address \(10\.0\.0\.8\)"):
convert_url_to_base64(url)
for verdict in _SSRF_VERDICTS:
assert attempts == [url]
def block(client, fetched_url, verdict=verdict, **kwargs):
attempts.append(fetched_url)
raise SSRFError(verdict)
monkeypatch.setattr(image_handling, "safe_get", block)
with pytest.raises(litellm.ImageFetchError) as raised:
convert_url_to_base64(url)
messages.append(raised.value.message)
assert attempts == [url] * len(_SSRF_VERDICTS)
_assert_one_verdict_free_message(messages, url)
async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch):

78
tests/test_litellm/llms/black_forest_labs/image_edit/test_bfl_image_edit_transformation.py
View file

@ -16,6 +16,9 @@ import httpx
import pytest
from litellm.llms.black_forest_labs.image_edit import (
transformation as bfl_transformation,
)
from litellm.llms.black_forest_labs.image_edit.transformation import (
BlackForestLabsImageEditConfig,
)
@ -186,7 +189,7 @@ class TestBlackForestLabsImageEditTransformation:
assert data["output_format"] == "jpeg"
# BFL uses JSON, not multipart - files should be empty
assert files == []
assert files == ()
def test_transform_image_edit_request_with_mask(self):
"""Test request transformation with mask for inpainting."""
@ -299,3 +302,76 @@ class TestBlackForestLabsImageEditTransformation:
def test_use_multipart_form_data_returns_false(self):
"""Test that use_multipart_form_data returns False for BFL."""
assert self.config.use_multipart_form_data() is False
async def test_async_transform_image_edit_request_downloads_url_images_with_the_async_fetcher(monkeypatch):
served = b"png-bytes-from-cdn"
fetched = []
def forbid_sync_fetch(client, url, **kwargs):
raise AssertionError(f"sync image fetch ran on the event loop: {url}")
async def serve(client, url, **kwargs):
fetched.append((url, kwargs.get("timeout")))
return httpx.Response(200, content=served, request=httpx.Request("GET", url))
monkeypatch.setattr(bfl_transformation, "safe_get", forbid_sync_fetch)
monkeypatch.setattr(bfl_transformation, "async_safe_get", serve)
data, files = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
model="flux-kontext-pro",
prompt="Add a red hat",
image="https://cdn.example/photo.png",
image_edit_optional_request_params={"mask": "https://cdn.example/mask.png", "seed": 7},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert base64.b64decode(data["input_image"]) == served
assert base64.b64decode(data["mask"]) == served
assert data["seed"] == 7
assert files == ()
assert fetched == [("https://cdn.example/photo.png", 60.0), ("https://cdn.example/mask.png", 60.0)]
async def test_async_transform_image_edit_request_never_fetches_for_local_images(monkeypatch):
def refuse(*args, **kwargs):
raise AssertionError("no network fetch expected for local image bytes")
monkeypatch.setattr(bfl_transformation, "safe_get", refuse)
monkeypatch.setattr(bfl_transformation, "async_safe_get", refuse)
data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
model="flux-kontext-pro",
prompt="Add a red hat",
image=[BytesIO(b"first"), BytesIO(b"other")],
image_edit_optional_request_params={"mask": b"mask-bytes"},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert base64.b64decode(data["input_image"]) == b"first"
assert base64.b64decode(data["mask"]) == b"mask-bytes"
async def test_async_transform_image_edit_request_downloads_only_the_first_url_of_a_list(monkeypatch):
fetched = []
async def serve(client, url, **kwargs):
fetched.append(url)
return httpx.Response(200, content=b"first-bytes", request=httpx.Request("GET", url))
monkeypatch.setattr(bfl_transformation, "safe_get", lambda *args, **kwargs: pytest.fail("sync fetch ran"))
monkeypatch.setattr(bfl_transformation, "async_safe_get", serve)
data, _ = await BlackForestLabsImageEditConfig().async_transform_image_edit_request(
model="flux-kontext-pro",
prompt="Add a red hat",
image=["https://cdn.example/a.png", "https://cdn.example/b.png"],
image_edit_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert fetched == ["https://cdn.example/a.png"]
assert base64.b64decode(data["input_image"]) == b"first-bytes"

163
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -19,6 +19,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import (
BaseAudioTranscriptionConfig,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
@ -31,7 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ModelResponse, TranscriptionResponse
from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse
_ACTIVE_KEY = "_code_interpreter_interception_active"
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
@ -3244,6 +3245,11 @@ class _TransformRecordingConfig(BaseConfig):
def get_error_class(self, error_message, status_code, headers):
return BaseLLMException(status_code=status_code, message=error_message, headers=headers)
def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False):
return litellm.OpenAIGPTConfig().get_model_response_iterator(
streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode
)
def _start_async_completion(config, logging_obj=None):
captured = {}
@ -3312,3 +3318,158 @@ async def test_completion_keeps_sync_transform_request_before_returning_by_defau
assert config.transform_calls == ["sync"]
assert captured["body"] == {"transformed_by": "sync"}
assert response.choices[0].message.content == "sync"
def _sse_echoing_transformed_by(request):
transformed_by = json.loads(request.content)["transformed_by"]
chunk = {
"id": "chatcmpl-1",
"object": "chat.completion.chunk",
"created": 1,
"model": "stub-model",
"choices": [{"index": 0, "delta": {"content": transformed_by}, "finish_reason": None}],
}
return httpx.Response(
200,
content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode(),
headers={"content-type": "text/event-stream"},
request=request,
)
def _streaming_logging_obj():
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj = Logging(
model="stub-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="async-transform-stream",
function_id="f",
)
logging_obj.update_environment_variables(
model="stub-model", user="", optional_params={}, litellm_params={}, custom_llm_provider="openai"
)
return logging_obj
async def test_completion_streams_after_the_async_transform_request():
config = _TransformRecordingConfig(transform_async=True)
loop_thread = threading.current_thread()
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by))
stream = await BaseLLMHTTPHandler().completion(
model="stub-model",
messages=[{"role": "user", "content": "hi"}],
api_base="https://llm.example/v1/chat",
custom_llm_provider="openai",
model_response=ModelResponse(),
encoding=None,
logging_obj=_streaming_logging_obj(),
optional_params={},
timeout=10.0,
litellm_params={},
acompletion=True,
stream=True,
client=client,
provider_config=config,
)
collected = [chunk async for chunk in stream]
assert config.transform_calls == ["async"]
assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads)
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async"
class _ImageEditRecordingConfig(BaseImageEditConfig):
def __init__(self):
self.transform_calls = []
def get_supported_openai_params(self, model):
return []
def map_openai_params(self, image_edit_optional_params, model, drop_params):
return dict(image_edit_optional_params)
def validate_environment(self, headers, model, api_key=None, litellm_params=None, api_base=None):
return {}
def get_complete_url(self, model, api_base, litellm_params):
return "https://images.example/v1/edits"
def use_multipart_form_data(self):
return False
def transform_image_edit_request(
self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers
):
self.transform_calls.append("sync")
return {"transformed_by": "sync"}, []
async def async_transform_image_edit_request(
self, model, prompt, image, image_edit_optional_request_params, litellm_params, headers
):
self.transform_calls.append("async")
return {"transformed_by": "async"}, []
def transform_image_edit_response(self, model, raw_response, logging_obj):
return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["transformed_by"])])
def _echo_json_transport(captured):
def handle(request):
captured["body"] = json.loads(request.content)
return httpx.Response(200, json=captured["body"])
return httpx.MockTransport(handle)
async def test_async_image_edit_handler_awaits_the_async_transform():
config = _ImageEditRecordingConfig()
captured = {}
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=_echo_json_transport(captured))
response = await BaseLLMHTTPHandler().async_image_edit_handler(
model="edit-model",
image=b"raw-image",
prompt="add a hat",
image_edit_provider_config=config,
image_edit_optional_request_params={},
custom_llm_provider="openai",
litellm_params=GenericLiteLLMParams(),
logging_obj=Mock(),
timeout=10.0,
client=client,
)
assert config.transform_calls == ["async"]
assert captured["body"] == {"transformed_by": "async"}
assert response.data[0].b64_json == "async"
def test_image_edit_handler_keeps_the_sync_transform():
config = _ImageEditRecordingConfig()
captured = {}
client = HTTPHandler()
client.client = httpx.Client(transport=_echo_json_transport(captured))
response = BaseLLMHTTPHandler().image_edit_handler(
model="edit-model",
image=b"raw-image",
prompt="add a hat",
image_edit_provider_config=config,
image_edit_optional_request_params={},
custom_llm_provider="openai",
litellm_params=GenericLiteLLMParams(),
logging_obj=Mock(),
timeout=10.0,
client=client,
)
assert config.transform_calls == ["sync"]
assert captured["body"] == {"transformed_by": "sync"}
assert response.data[0].b64_json == "sync"