diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index 6be8d0683..81285788a 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -387,11 +387,14 @@ async def generate_from_chat_context( Returns: ModelOutputThunk[C]: A thunk holding the (lazy) model output. + Ollama requires base64-encoded images, so any `ImageUrlBlock` in a + message is fetched and encoded automatically before the request. + Raises: RuntimeError: If not called from a thread with a running event loop. - ValueError: If a message contains an ``ImageUrlBlock``; Ollama requires - base64-encoded images — convert to an ``ImageBlock`` first. - ValueError: If a message contains an ``AudioBlock`` or ``AudioUrlBlock``; + ValueError: If a message contains an `ImageUrlBlock` whose image + cannot be downloaded or decoded. + ValueError: If a message contains an `AudioBlock` or `AudioUrlBlock`; Ollama does not support audio input. """ # Start by awaiting any necessary computation. @@ -425,13 +428,25 @@ async def generate_from_chat_context( # to print `Message`s to correctly serialize any documents with the message. Do the printing here. replay_flags = should_replay_reasoning(messages, self._provider) for m, replay in zip(messages, replay_flags): + image_values: list[str] | None = None if m.images is not None: - for img in m.images: - if isinstance(img, ImageUrlBlock): - raise ValueError( - "OllamaModelBackend does not support URL images (ImageUrlBlock). " - "Convert the image to a base64-encoded ImageBlock before passing it to Ollama." - ) + # Ollama only accepts base64-encoded images, so URL images are + # downloaded and encoded on the fly rather than rejected. + # `resolve_base64` memoizes on the block, so re-using the same + # block across turns downloads only once. The download is + # blocking, so offload each one to a thread and run them + # concurrently to avoid stalling the event loop; non-URL images + # already carry their base64 value. + image_values = [str(img.value) for img in m.images] + url_downloads = { + i: asyncio.to_thread(img.resolve_base64) + for i, img in enumerate(m.images) + if isinstance(img, ImageUrlBlock) + } + if url_downloads: + downloaded = await asyncio.gather(*url_downloads.values()) + for i, value in zip(url_downloads.keys(), downloaded): + image_values[i] = value if m.audio: raise ValueError( "OllamaModelBackend does not support audio (AudioBlock/AudioUrlBlock). " @@ -441,9 +456,7 @@ async def generate_from_chat_context( "role": m.role, "content": self.formatter.print(m), "images": ( - _strip_data_uri_prefix([str(img.value) for img in m.images]) - if m.images - else None + _strip_data_uri_prefix(image_values) if image_values else None ), } # Ollama's native SDK carries reasoning under the `thinking` key (see diff --git a/mellea/core/__init__.py b/mellea/core/__init__.py index c1fba8669..62e2c472f 100644 --- a/mellea/core/__init__.py +++ b/mellea/core/__init__.py @@ -32,6 +32,7 @@ TemplateRepresentation, blockify, get_audio_from_component, + make_image_block, ) from .formatter import Formatter from .requirement import ( @@ -93,5 +94,6 @@ def __getattr__(name: str) -> object: "generate_walk", "get_audio_from_component", "log_context", + "make_image_block", "set_log_context", ] diff --git a/mellea/core/base.py b/mellea/core/base.py index d36bb4d08..89cfc378f 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -20,6 +20,8 @@ import datetime import enum import logging +import threading +from collections import OrderedDict from collections.abc import Callable, Coroutine, Iterable, Mapping from copy import copy, deepcopy from dataclasses import dataclass, field @@ -38,6 +40,7 @@ if TYPE_CHECKING: import torch +import requests import typing_extensions from PIL import Image as PILImage @@ -193,8 +196,8 @@ class ImageUrlBlock(CBlock): Use this when the image is hosted remotely and you want to pass the URL directly to backends that support it (e.g. OpenAI). Backends that only - accept base64-encoded images (e.g. Ollama) will raise a ``ValueError`` - rather than silently drop the image. + accept base64-encoded images (e.g. Ollama) download and encode the image + automatically, so the URL never has to be converted by hand. Args: value (str): A URL string pointing to the image. @@ -206,8 +209,8 @@ def __init__(self, value: str, meta: dict[str, Any] | None = None): """Initialize ImageUrlBlock with a URL string. Raises: - ValueError: If ``value`` does not look like a URL (does not start - with ``http://`` or ``https://``). + ValueError: If `value` does not look like a URL (does not start + with `http://` or `https://`). """ if not value.startswith(("http://", "https://")): raise ValueError( @@ -215,6 +218,25 @@ def __init__(self, value: str, meta: dict[str, Any] | None = None): ) super().__init__(value, meta) + def resolve_base64(self) -> str: + """Return the image as a base64-encoded PNG, downloading it once per URL. + + Backends that cannot pass a URL through (e.g. Ollama) need the image + as base64. The download and encode result is memoized in a process-wide, + URL-keyed cache, so re-using the URL across conversation turns — even + via a freshly reconstructed block — does not re-fetch the image. This + call is blocking; async callers should offload it with + `asyncio.to_thread`. + + Returns: + str: The base64-encoded PNG representation of the image at the URL. + + Raises: + ValueError: If the response exceeds the size cap or the image + cannot be downloaded or decoded. + """ + return _cached_download_image_as_base64(str(self.value)) + def __repr__(self) -> str: """Provides a python-parsable representation of the block (usually).""" return f"ImageUrlBlock({self.value}, {self._meta.__repr__()})" @@ -342,6 +364,161 @@ def __repr__(self) -> str: return f"AudioUrlBlock({self.value}, {self.format}, {self._meta.__repr__()})" +_IMAGE_DOWNLOAD_TIMEOUT_S: float = 10.0 +"""Socket timeout (seconds) applied to image URL downloads.""" + +_IMAGE_DOWNLOAD_MAX_BYTES: int = 25 * 1024 * 1024 +"""Maximum accepted size (bytes) of a downloaded image body.""" + +_IMAGE_CACHE_MAX_ENTRIES: int = 128 +"""Maximum number of URL -> base64 entries retained by the download cache.""" + +_image_base64_cache: OrderedDict[str, str] = OrderedDict() +"""Process-wide LRU cache mapping image URLs to their base64-encoded PNG.""" + +_image_base64_cache_lock = threading.Lock() +"""Guards `_image_base64_cache` against concurrent `asyncio.to_thread` callers.""" + + +def _cached_download_image_as_base64(url: str) -> str: + """Download an image as base64, memoizing the result per URL. + + Wraps `_download_image_as_base64` with a bounded, thread-safe LRU cache + keyed on the URL so the same image is fetched only once regardless of how + many `ImageUrlBlock` instances reference it. The download runs outside the + lock so concurrent fetches of distinct URLs still proceed in parallel; + only the small cache read/write is serialized. + + Args: + url: An `http://` or `https://` URL pointing to an image. + + Returns: + str: The base64-encoded PNG representation of the image at the URL. + + Raises: + ValueError: If the response exceeds the size cap or the image cannot + be downloaded or decoded. + """ + with _image_base64_cache_lock: + cached = _image_base64_cache.get(url) + if cached is not None: + _image_base64_cache.move_to_end(url) # mark as most-recently used + return cached + + # Download outside the lock; two callers racing on a cold URL may both + # fetch, but the result is identical and the last writer simply wins. + encoded = _download_image_as_base64(url) + + with _image_base64_cache_lock: + _image_base64_cache[url] = encoded + _image_base64_cache.move_to_end(url) + while len(_image_base64_cache) > _IMAGE_CACHE_MAX_ENTRIES: + _image_base64_cache.popitem(last=False) # evict least-recently used + return encoded + + +def _download_image_as_base64(url: str) -> str: + """Download an image from a URL and return it as a base64-encoded PNG string. + + Fetches the bytes at `url`, loads them through PIL to confirm they are a + real image, and re-encodes the result as a base64 PNG so the output is + consistent with `ImageBlock`'s expected format. + + A timeout bounds slow responses and the body is streamed with a size cap + to guard against memory-exhaustion. This function is blocking; async + callers should offload it with `asyncio.to_thread`. + + Args: + url: An `http://` or `https://` URL pointing to an image. + + Returns: + str: The base64-encoded PNG representation of the downloaded image. + + Raises: + ValueError: If the response exceeds the size cap or the image cannot + be downloaded or decoded. + """ + try: + with requests.get( # scheme validated by caller + url, timeout=_IMAGE_DOWNLOAD_TIMEOUT_S, stream=True + ) as response: + response.raise_for_status() + declared = response.headers.get("Content-Length") + if declared is not None and int(declared) > _IMAGE_DOWNLOAD_MAX_BYTES: + raise ValueError( + f"Image at {url!r} exceeds the {_IMAGE_DOWNLOAD_MAX_BYTES}-byte limit" + ) + # Stream so an undeclared/lying Content-Length can't exhaust memory. + raw = response.raw.read(_IMAGE_DOWNLOAD_MAX_BYTES + 1, decode_content=True) + if len(raw) > _IMAGE_DOWNLOAD_MAX_BYTES: + raise ValueError( + f"Image at {url!r} exceeds the {_IMAGE_DOWNLOAD_MAX_BYTES}-byte limit" + ) + image = PILImage.open(BytesIO(raw)) + except (requests.RequestException, OSError, ValueError) as e: + raise ValueError( + f"Failed to download or decode image from URL {url!r}: {e}" + ) from e + return ImageBlock.pil_to_base64(image) + + +def make_image_block( + src: str | PILImage.Image, + *, + convert_to_base64: bool = False, + meta: dict[str, Any] | None = None, +) -> ImageBlock | ImageUrlBlock: + """Create the appropriate image block from any supported image source. + + Dispatches on the type and shape of `src` so callers don't have to know + whether they need an `ImageBlock` (base64-encoded) or an `ImageUrlBlock` + (URL-referenced): + + - A PIL image is encoded to a base64 PNG and returned as an `ImageBlock`. + - An `http://`/`https://` URL is returned as an `ImageUrlBlock`, unless + `convert_to_base64=True`, in which case the image is downloaded and + returned as an `ImageBlock`. + - A base64-encoded PNG string (with or without a data URI prefix) is + returned as an `ImageBlock`. + + Args: + src: The image source — a PIL image, an image URL, or a base64-encoded + PNG string. + convert_to_base64: If `True` and `src` is a URL, download the image and + return an `ImageBlock` instead of an `ImageUrlBlock`. Ignored for + non-URL sources. + meta: Optional metadata to associate with the returned block. + + Returns: + ImageBlock | ImageUrlBlock: An `ImageBlock` for PIL images, base64 + strings, and downloaded URLs; an `ImageUrlBlock` for URLs when + `convert_to_base64` is `False`. + + Raises: + ValueError: If `src` is a string that is neither a valid URL nor a + valid base64-encoded PNG, or if a URL download fails. + TypeError: If `src` is not a PIL image or a string. + """ + if isinstance(src, PILImage.Image): + return ImageBlock.from_pil_image(src, meta) + + if isinstance(src, str): + if src.startswith(("http://", "https://")): + if convert_to_base64: + return ImageBlock(_download_image_as_base64(src), meta) + return ImageUrlBlock(src, meta) + if ImageBlock.is_valid_base64_png(src): + return ImageBlock(src, meta) + raise ValueError( + f"make_image_block could not interpret string source; expected an " + f"http(s) URL or a base64-encoded PNG, got: {src!r}" + ) + + raise TypeError( + f"make_image_block expects a PIL image or a string source, got: {type(src)!r}" + ) + + S = typing_extensions.TypeVar("S", default=Any, covariant=True) """Used for class definitions for Component and ModelOutputThunk; also used for functions that don't accept CBlocks. Defaults to `Any`.""" @@ -1623,9 +1800,9 @@ def get_images_from_component(c: Component) -> None | list[ImageBlock | ImageUrl c: The `Component` whose `images` attribute is inspected. Returns: - A non-empty list of ``ImageBlock`` or ``ImageUrlBlock`` objects if the - component has an ``images`` attribute with at least one element; - ``None`` otherwise. + A non-empty list of `ImageBlock` or `ImageUrlBlock` objects if the + component has an `images` attribute with at least one element; + `None` otherwise. """ if hasattr(c, "images"): imgs = c.images # type: ignore diff --git a/mellea/stdlib/components/chat.py b/mellea/stdlib/components/chat.py index d3b3eb2cb..4caa62c2c 100644 --- a/mellea/stdlib/components/chat.py +++ b/mellea/stdlib/components/chat.py @@ -40,8 +40,9 @@ class Message(Component["Message"]): images (list[ImageBlock | ImageUrlBlock] | None): Optional images associated with the message. Use `ImageBlock` for base64-encoded images (supported by all vision backends) or `ImageUrlBlock` for - URL-referenced images (supported by OpenAI-compatible backends only; - backends that require base64 will raise a ``ValueError``). + URL-referenced images (passed directly to OpenAI-compatible + backends; backends that require base64, such as Ollama, download + and encode the image automatically). audio (list[AudioBlock | AudioUrlBlock] | None): Optional audio associated with the message. documents (list[Document] | None): Optional documents associated with diff --git a/test/backends/test_vision_ollama.py b/test/backends/test_vision_ollama.py index b03cd0956..d1ecd06f2 100644 --- a/test/backends/test_vision_ollama.py +++ b/test/backends/test_vision_ollama.py @@ -45,6 +45,16 @@ # ── Shared image fixture ────────────────────────────────────────────────────── +@pytest.fixture(autouse=True) +def _clear_image_cache(): + """Isolate tests from the process-wide URL -> base64 download cache.""" + from mellea.core import base as base_mod + + base_mod._image_base64_cache.clear() + yield + base_mod._image_base64_cache.clear() + + @pytest.fixture(scope="module") def pil_image(): rng = np.random.default_rng(seed=42) @@ -138,12 +148,45 @@ def test_image_block_in_instruction( assert image_list[0] == str(image_block) -def test_image_url_block_rejected_by_ollama(mocked_session: MelleaSession): +def test_image_url_block_auto_downloaded_by_ollama( + mocked_session: MelleaSession, pil_image: Image.Image +): + # Ollama only accepts base64 images, so a URL image must be downloaded and + # encoded transparently rather than rejected. + encoded = ImageBlock.from_pil_image(pil_image).value url_block: ImageUrlBlock = ImageUrlBlock("https://example.com/photo.png") images: list[ImageBlock | ImageUrlBlock] = [url_block] - with pytest.raises(ValueError, match="ImageUrlBlock"): + + with patch( + "mellea.core.base._download_image_as_base64", return_value=encoded + ) as mock_download: mocked_session.chat("What is in this image?", images=images) + mock_download.assert_called_once_with("https://example.com/photo.png") + + turn = mocked_session.ctx.last_turn() + assert turn is not None + lp = turn.output._generate_log.prompt # type: ignore[union-attr] + assert isinstance(lp, list) + prompt_msg = lp[0] + image_list = prompt_msg.get("images") + assert isinstance(image_list, list) + assert len(image_list) == 1 + # The downloaded base64 (data-URI-stripped) is embedded in the payload. + assert image_list[0] == encoded + + +def test_image_url_block_download_failure_raises(mocked_session: MelleaSession): + url_block: ImageUrlBlock = ImageUrlBlock("https://example.com/photo.png") + images: list[ImageBlock | ImageUrlBlock] = [url_block] + + with patch( + "mellea.core.base._download_image_as_base64", + side_effect=ValueError("Failed to download or decode image from URL"), + ): + with pytest.raises(ValueError, match="Failed to download"): + mocked_session.chat("What is in this image?", images=images) + def test_audio_block_rejected_by_ollama(mocked_session: MelleaSession): """AudioBlock raises ValueError — Ollama does not support audio input.""" @@ -172,6 +215,53 @@ def test_audio_url_block_rejected_by_ollama(mocked_session: MelleaSession): mocked_session.chat("Transcribe this.", audio=[url_block]) +def test_image_url_block_drives_real_download( + mocked_session: MelleaSession, pil_image: Image.Image +): + """Exercise the real `_download_image_as_base64` through the Ollama path. + + Patches only the network layer (`requests.get`) so the block's + `resolve_base64` call, thread-offload, and payload wiring are all covered + — this catches regressions the fully-mocked download tests cannot. + """ + buf = BytesIO() + pil_image.save(buf, format="PNG") + png_bytes = buf.getvalue() + + url_block: ImageUrlBlock = ImageUrlBlock("https://example.com/photo.png") + images: list[ImageBlock | ImageUrlBlock] = [url_block] + + class _FakeRaw: + def read(self, amt=None, decode_content=False): + return png_bytes if amt is None else png_bytes[:amt] + + class _FakeResponse: + headers: dict[str, str] = {} + raw = _FakeRaw() + + def raise_for_status(self): + return None + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + with patch("mellea.core.base.requests.get", return_value=_FakeResponse()): + mocked_session.chat("What is in this image?", images=images) + + turn = mocked_session.ctx.last_turn() + assert turn is not None + lp = turn.output._generate_log.prompt # type: ignore[union-attr] + assert isinstance(lp, list) + image_list = lp[0].get("images") + assert isinstance(image_list, list) + assert len(image_list) == 1 + # The real helper re-encodes as base64 PNG (data-URI-stripped in the payload). + assert base64.b64decode(image_list[0]) + + def test_image_block_in_chat(mocked_session: MelleaSession, pil_image: Image.Image): image_block = ImageBlock.from_pil_image(pil_image) ct = mocked_session.chat( diff --git a/test/core/test_base.py b/test/core/test_base.py index 1cd9d9cb6..a5a0d2c4d 100644 --- a/test/core/test_base.py +++ b/test/core/test_base.py @@ -18,6 +18,7 @@ RawProviderResponse, blockify, get_audio_from_component, + make_image_block, ) from mellea.core.backend import generate_walk from mellea.stdlib.components import Message @@ -332,6 +333,188 @@ def _parse(self, computed: ModelOutputThunk) -> str: assert get_audio_from_component(_ComponentWithoutAudio()) is None +# --- make_image_block factory --- + + +def _png_bytes() -> bytes: + img = PILImage.new("RGB", (1, 1), color="blue") + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +def test_make_image_block_from_pil(): + img = PILImage.new("RGB", (1, 1), color="green") + block = make_image_block(img) + assert isinstance(block, ImageBlock) + assert ImageBlock.is_valid_base64_png(str(block)) + + +def test_make_image_block_from_url_returns_url_block(): + block = make_image_block("https://example.com/cat.png") + assert isinstance(block, ImageUrlBlock) + assert block.value == "https://example.com/cat.png" + + +def test_make_image_block_from_base64_returns_image_block(): + b64 = _make_png_b64() + block = make_image_block(b64) + assert isinstance(block, ImageBlock) + assert block.value == b64 + + +def test_make_image_block_data_uri_returns_image_block(): + data_uri = f"data:image/png;base64,{_make_png_b64()}" + block = make_image_block(data_uri) + assert isinstance(block, ImageBlock) + + +def test_make_image_block_preserves_meta(): + b64 = _make_png_b64() + block = make_image_block(b64, meta={"alt": "a dot"}) + assert block._meta == {"alt": "a dot"} + + +class _FakeRaw: + """Stand-in for `requests.Response.raw` supporting a capped `.read()`.""" + + def __init__(self, body: bytes): + self._body = body + + def read(self, amt: int | None = None, decode_content: bool = False) -> bytes: + return self._body if amt is None else self._body[:amt] + + +class _FakeResponse: + """Minimal stand-in for a `requests.Response` used as a context manager.""" + + def __init__(self, body: bytes, content_length: str | None = None): + self.raw = _FakeRaw(body) + self.headers = ( + {} if content_length is None else {"Content-Length": content_length} + ) + + def raise_for_status(self): + return None + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +@pytest.fixture(autouse=True) +def _clear_image_cache(): + """Isolate tests from the process-wide URL -> base64 download cache.""" + import mellea.core.base as base_mod + + base_mod._image_base64_cache.clear() + yield + base_mod._image_base64_cache.clear() + + +def _patch_download(monkeypatch, fake_get): + """Stub `requests.get` for a download test.""" + import mellea.core.base as base_mod + + monkeypatch.setattr(base_mod.requests, "get", fake_get) + + +def test_make_image_block_url_convert_to_base64(monkeypatch): + _patch_download(monkeypatch, lambda url, **kw: _FakeResponse(_png_bytes())) + + block = make_image_block("https://example.com/cat.png", convert_to_base64=True) + assert isinstance(block, ImageBlock) + assert ImageBlock.is_valid_base64_png(str(block)) + + +def test_make_image_block_download_failure_raises(monkeypatch): + import mellea.core.base as base_mod + + def boom(url, **kw): + raise base_mod.requests.RequestException("no network") + + _patch_download(monkeypatch, boom) + + with pytest.raises(ValueError, match="Failed to download"): + make_image_block("https://example.com/cat.png", convert_to_base64=True) + + +def test_make_image_block_download_rejects_oversized_body(monkeypatch): + from mellea.core import base as base_mod + + big = _png_bytes() * (base_mod._IMAGE_DOWNLOAD_MAX_BYTES + 1) + _patch_download(monkeypatch, lambda url, **kw: _FakeResponse(big)) + + with pytest.raises(ValueError, match="exceeds"): + make_image_block("https://example.com/big.png", convert_to_base64=True) + + +def test_make_image_block_download_rejects_oversized_content_length(monkeypatch): + from mellea.core import base as base_mod + + limit = base_mod._IMAGE_DOWNLOAD_MAX_BYTES + _patch_download( + monkeypatch, + lambda url, **kw: _FakeResponse(_png_bytes(), content_length=str(limit + 1)), + ) + + with pytest.raises(ValueError, match="exceeds"): + make_image_block("https://example.com/big.png", convert_to_base64=True) + + +def test_image_url_block_resolve_base64_caches_by_url(monkeypatch): + """resolve_base64 downloads once per URL; a reconstructed block hits cache.""" + calls: dict[str, int] = {} + + def fake_get(url, **kw): + calls[url] = calls.get(url, 0) + 1 + return _FakeResponse(_png_bytes()) + + _patch_download(monkeypatch, fake_get) + + url = "https://example.com/cat.png" + first = ImageUrlBlock(url).resolve_base64() + # A freshly reconstructed block for the same URL must not re-download. + second = ImageUrlBlock(url).resolve_base64() + + assert ImageBlock.is_valid_base64_png(first) + assert first == second + assert calls[url] == 1 # cache is keyed on the URL, not the instance + + # A different URL is a cache miss and triggers its own download. + other = "https://example.com/dog.png" + ImageUrlBlock(other).resolve_base64() + assert calls[other] == 1 + + +def test_image_cache_evicts_least_recently_used(monkeypatch): + """The download cache is bounded and evicts the least-recently-used URL.""" + from mellea.core import base as base_mod + + _patch_download(monkeypatch, lambda url, **kw: _FakeResponse(_png_bytes())) + monkeypatch.setattr(base_mod, "_IMAGE_CACHE_MAX_ENTRIES", 2) + + a, b, c = (f"https://example.com/{n}.png" for n in ("a", "b", "c")) + base_mod._cached_download_image_as_base64(a) + base_mod._cached_download_image_as_base64(b) + base_mod._cached_download_image_as_base64(a) # touch `a` so `b` is now LRU + base_mod._cached_download_image_as_base64(c) # evicts `b` + + assert set(base_mod._image_base64_cache) == {a, c} + + +def test_make_image_block_invalid_string_raises(): + with pytest.raises(ValueError, match="could not interpret"): + make_image_block("not-a-url-or-base64!!!") + + +def test_make_image_block_invalid_type_raises(): + with pytest.raises(TypeError, match="expects a PIL image or a string"): + make_image_block(12345) # type: ignore[arg-type] + + # --- ModelOutputThunk._copy_from ---