diff --git a/litellm/images/main.py b/litellm/images/main.py index a40a590799f..65efc45f4fb 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -176,7 +176,7 @@ def _filename_for_file(file_value: Any, default: str) -> str: def _rust_image_file_part(file_value: Any, default_filename: str) -> dict[str, object]: - filename = default_filename + filename: str | None = None content_type: str | None = None raw_file = file_value 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 c365f74ac09..a78f4b59416 100644 --- a/tests/test_litellm/images/test_rust_image_edit_bridge.py +++ b/tests/test_litellm/images/test_rust_image_edit_bridge.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib +from io import BytesIO from pathlib import Path import httpx @@ -85,6 +86,38 @@ class RecordingAsyncImageEditBridge: return {"created": 1, "data": [{"b64_json": "YXN5bmMtaW1hZ2U="}]} +class TextFileLike: + def __init__(self, data: str) -> None: + self.data = data + self.position = 0 + + def tell(self) -> int: + return self.position + + def seek(self, position: int) -> None: + self.position = position + + def read(self, *_args: object, **_kwargs: object) -> str: + self.position = len(self.data) + return self.data + + +class BytearrayFileLike: + def __init__(self, data: bytes) -> None: + self.data = bytearray(data) + self.position = 0 + + def tell(self) -> int: + return self.position + + def seek(self, position: int) -> None: + self.position = position + + def read(self, *_args: object, **_kwargs: object) -> bytearray: + self.position = len(self.data) + return self.data + + @pytest.fixture(autouse=True) def reset_rust_bridge() -> None: rust_bridge.use_litellm_rust( @@ -187,3 +220,87 @@ 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 + + +def test_timeout_to_seconds_returns_none_when_all_timeout_values_unset() -> None: + timeout = httpx.Timeout(connect=None, read=None, write=None, pool=None) + + assert image_main._timeout_to_seconds(timeout) is None + + +def test_rust_image_file_part_reads_bytearray() -> None: + part = image_main._rust_image_file_part(bytearray(PNG_BYTES), "default.png") + + assert part["filename"] == "default.png" + assert part["content_type"] == "image/png" + assert part["data_base64"] == "iVBORw0KGgpmYWtlLWltYWdl" + + +def test_rust_image_file_part_preserves_bytesio_position() -> None: + image = BytesIO(PNG_BYTES) + image.seek(4) + + part = image_main._rust_image_file_part(image, "buffer.png") + + assert image.tell() == 4 + assert part["filename"] == "buffer.png" + assert part["content_type"] == "image/png" + assert part["data_base64"] == "iVBORw0KGgpmYWtlLWltYWdl" + + +def test_rust_image_file_part_uses_named_file_object_filename() -> None: + image = BytesIO(PNG_BYTES) + image.name = "/tmp/named-source.png" # type: ignore[attr-defined] + + part = image_main._rust_image_file_part(image, "buffer.png") + + assert part["filename"] == "named-source.png" + assert part["data_base64"] == "iVBORw0KGgpmYWtlLWltYWdl" + + +def test_rust_image_file_part_reads_text_file_like() -> None: + image = TextFileLike("plain-text") + image.seek(5) + + part = image_main._rust_image_file_part(image, "text.png") + + assert image.tell() == 5 + assert part["filename"] == "text.png" + assert part["data_base64"] == "cGxhaW4tdGV4dA==" + + +def test_rust_image_file_part_reads_bytearray_file_like() -> None: + image = BytearrayFileLike(PNG_BYTES) + image.seek(3) + + part = image_main._rust_image_file_part(image, "bytes.png") + + assert image.tell() == 3 + assert part["filename"] == "bytes.png" + assert part["data_base64"] == "iVBORw0KGgpmYWtlLWltYWdl" + + +def test_rust_image_file_part_reads_filesystem_path(tmp_path: Path) -> None: + image_path = tmp_path / "source.png" + image_path.write_bytes(PNG_BYTES) + + part = image_main._rust_image_file_part(str(image_path), "fallback.png") + + assert part["filename"] == "source.png" + assert part["data_base64"] == "iVBORw0KGgpmYWtlLWltYWdl" + + +def test_rust_image_file_part_reads_tuple_metadata() -> None: + part = image_main._rust_image_file_part( + ("custom.jpeg", BytesIO(PNG_BYTES), "image/jpeg"), + "fallback.png", + ) + + assert part["filename"] == "custom.jpeg" + assert part["content_type"] == "image/jpeg" + assert part["data_base64"] == "iVBORw0KGgpmYWtlLWltYWdl" + + +def test_rust_image_file_part_rejects_unsupported_input() -> None: + with pytest.raises(TypeError, match="Unsupported image file type"): + image_main._rust_image_file_part(object(), "fallback.png") diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7bbada1c5aa..b4abbf2bdd8 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -296,6 +296,8 @@ def test_load_rust_ocr_none_when_extension_absent(): litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI assert rust_bridge.load_rust_ocr() is None assert rust_bridge.load_rust_aocr() is None + assert rust_bridge.load_rust_image_edit() is None + assert rust_bridge.load_rust_aimage_edit() is None def test_load_rust_ocr_uses_compiled_extension(monkeypatch): @@ -305,11 +307,15 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): fake_module = types.ModuleType("litellm_python_bridge") fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] + fake_module.image_edit = lambda **kwargs: {"data": []} # type: ignore[attr-defined] + fake_module.aimage_edit = lambda **kwargs: {"data": []} # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension assert rust_bridge.load_rust_ocr() is fake_module.ocr assert rust_bridge.load_rust_aocr() is fake_module.aocr + assert rust_bridge.load_rust_image_edit() is fake_module.image_edit + assert rust_bridge.load_rust_aimage_edit() is fake_module.aimage_edit def test_timeout_to_seconds_handles_float_timeout_and_none():