diff --git a/litellm/images/main.py b/litellm/images/main.py index 4da8edd25b4..a40a590799f 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -81,6 +81,7 @@ from litellm.ocr.rust_bridge import ( RustImageEdit, load_rust_aimage_edit, load_rust_image_edit, + rust_image_edit_enabled, ) # Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls @@ -119,10 +120,15 @@ def _timeout_to_seconds( if timeout is None: return None if isinstance(timeout, httpx.Timeout): - read_timeout = timeout.read - if read_timeout is None: - return None - return float(read_timeout) + for timeout_value in ( + timeout.read, + timeout.connect, + timeout.write, + timeout.pool, + ): + if timeout_value is not None: + return float(timeout_value) + return None return float(timeout) @@ -147,11 +153,14 @@ def _read_file_content(file_value: Any) -> bytes: if isinstance(data, str): return data.encode() return bytes(data) - if isinstance(file_value, (str, os.PathLike)) and os.path.exists(file_value): - with open(file_value, "rb") as file: + if isinstance(file_value, (str, os.PathLike)): + path = os.fspath(file_value) + if not os.path.exists(path): + raise FileNotFoundError( + f"Image file path does not exist for Rust image_edit: {path}" + ) + with open(path, "rb") as file: return file.read() - if isinstance(file_value, str): - return file_value.encode() raise TypeError( f"Unsupported image file type for Rust image_edit: {type(file_value)}" ) @@ -1065,7 +1074,7 @@ def image_edit( ) local_vars.update(kwargs) - if custom_llm_provider == "vllm": + if custom_llm_provider == "vllm" and rust_image_edit_enabled(): image_edit_optional_params = ( _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( local_vars diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index b7f59402aaf..bee36318878 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -88,6 +88,7 @@ class _Unset: _UNSET: Final[_Unset] = _Unset() _rust_ocr_enabled = False +_rust_image_edit_enabled = False _rust_ocr_impl: RustOcr | None = None _rust_aocr_impl: RustAocr | None = None _rust_image_edit_impl: RustImageEdit | None = None @@ -108,9 +109,10 @@ def use_litellm_rust( extension is loaded on demand and any previously injected bridge is preserved. Pass ``None`` explicitly to clear a prior injection. """ - global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl + global _rust_ocr_enabled, _rust_image_edit_enabled, _rust_ocr_impl, _rust_aocr_impl global _rust_image_edit_impl, _rust_aimage_edit_impl _rust_ocr_enabled = enabled + _rust_image_edit_enabled = enabled if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr if not isinstance(aocr, _Unset): @@ -128,7 +130,7 @@ def rust_ocr_enabled() -> bool: def rust_image_edit_enabled() -> bool: """Whether the Rust image-edit path has been turned on.""" - return _rust_ocr_enabled + return _rust_image_edit_enabled def load_rust_ocr() -> RustOcr | None: diff --git a/tests/test_litellm/images/test_rust_image_edit_bridge.py b/tests/test_litellm/images/test_rust_image_edit_bridge.py index a2615975616..c365f74ac09 100644 --- a/tests/test_litellm/images/test_rust_image_edit_bridge.py +++ b/tests/test_litellm/images/test_rust_image_edit_bridge.py @@ -3,10 +3,14 @@ from __future__ import annotations import importlib +from pathlib import Path +import httpx import pytest import litellm +from litellm.exceptions import APIConnectionError +from litellm.images import main as image_main rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") @@ -142,3 +146,44 @@ async def test_vllm_aimage_edit_routes_to_async_rust_bridge() -> None: assert len(bridge.calls) == 1 assert bridge.calls[0]["model"] == "qwen-image-edit" assert bridge.calls[0]["prompt"] == "make it darker" + + +def test_vllm_image_edit_respects_disabled_rust_bridge() -> None: + bridge = RecordingImageEditBridge() + litellm.use_litellm_rust(False, image_edit=bridge) + + with pytest.raises( + APIConnectionError, match="image edit is not supported for vllm" + ): + litellm.image_edit( + model="vllm/qwen-image-edit", + image=PNG_BYTES, + prompt="make it brighter", + api_base="http://localhost:8000", + ) + + assert bridge.calls == [] + + +def test_vllm_image_edit_missing_path_raises_before_bridge_call( + tmp_path: Path, +) -> None: + bridge = RecordingImageEditBridge() + missing_image = tmp_path / "missing-image.png" + litellm.use_litellm_rust(True, image_edit=bridge) + + with pytest.raises(APIConnectionError, match="Image file path does not exist"): + litellm.image_edit( + model="vllm/qwen-image-edit", + image=str(missing_image), + prompt="make it brighter", + api_base="http://localhost:8000", + ) + + assert bridge.calls == [] + + +def test_timeout_to_seconds_falls_back_to_non_read_timeout() -> None: + timeout = httpx.Timeout(connect=7.0, read=None, write=8.0, pool=9.0) + + assert image_main._timeout_to_seconds(timeout) == 7.0 diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index ccb11c0f0c6..7bbada1c5aa 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -240,10 +240,13 @@ def fake_async_bridge(): def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False + assert rust_bridge.rust_image_edit_enabled() is False litellm.use_litellm_rust() assert rust_bridge.rust_ocr_enabled() is True + assert rust_bridge.rust_image_edit_enabled() is True litellm.use_litellm_rust(False) assert rust_bridge.rust_ocr_enabled() is False + assert rust_bridge.rust_image_edit_enabled() is False def test_load_rust_ocr_returns_injected_impl():