From 74886082eb93324e6e43461e63808bb1eb1bf6c9 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Fri, 10 Jul 2026 14:47:47 -0400 Subject: [PATCH 1/3] feat: imageblock and imgurblock ergonomics Signed-off-by: AngeloDanducci --- mellea/backends/ollama.py | 26 ++++---- mellea/core/__init__.py | 2 + mellea/core/base.py | 99 +++++++++++++++++++++++++++-- mellea/stdlib/components/chat.py | 5 +- test/backends/test_vision_ollama.py | 37 ++++++++++- test/core/test_base.py | 87 +++++++++++++++++++++++++ 6 files changed, 233 insertions(+), 23 deletions(-) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index c6bad8317..3227d6a21 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -25,7 +25,7 @@ ModelToolCall, RawProviderResponse, ) -from ..core.base import AbstractMelleaTool +from ..core.base import AbstractMelleaTool, _download_image_as_base64 from ..formatters import ChatFormatter, TemplateFormatter from ..helpers import ( DEFAULT_CHUNK_TIMEOUT, @@ -388,8 +388,9 @@ async def generate_from_chat_context( 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 `ImageUrlBlock` whose image + cannot be downloaded or decoded; Ollama requires base64-encoded + images, so URL images are fetched and encoded automatically. """ # Start by awaiting any necessary computation. await self.do_generate_walk(action) @@ -421,21 +422,22 @@ async def generate_from_chat_context( # NOTE: `self.formatter.to_chat_messages` explicitly skips `Message` objects. However, we need # to print `Message`s to correctly serialize any documents with the message. Do the printing here. for m in messages: + 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. + image_values = [ + _download_image_as_base64(str(img.value)) + if isinstance(img, ImageUrlBlock) + else str(img.value) + for img in m.images + ] conversation.append( { "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 ), } ) diff --git a/mellea/core/__init__.py b/mellea/core/__init__.py index 58c7cde5a..c42c21063 100644 --- a/mellea/core/__init__.py +++ b/mellea/core/__init__.py @@ -29,6 +29,7 @@ S, TemplateRepresentation, blockify, + make_image_block, ) from .formatter import Formatter from .requirement import ( @@ -87,5 +88,6 @@ def __getattr__(name: str) -> object: "default_output_to_bool", "generate_walk", "log_context", + "make_image_block", "set_log_context", ] diff --git a/mellea/core/base.py b/mellea/core/base.py index f2655504d..514ef221b 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -20,6 +20,8 @@ import datetime import enum import logging +import urllib.error +import urllib.request from collections.abc import Callable, Coroutine, Iterable, Mapping from copy import copy, deepcopy from dataclasses import dataclass, field @@ -193,8 +195,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 +208,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( @@ -220,6 +222,89 @@ def __repr__(self) -> str: return f"ImageUrlBlock({self.value}, {self._meta.__repr__()})" +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. + + 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 image cannot be downloaded or the downloaded bytes + cannot be decoded as an image. + """ + try: + with urllib.request.urlopen(url) as response: # scheme validated by caller + raw = response.read() + image = PILImage.open(BytesIO(raw)) + except (urllib.error.URLError, OSError, ValueError) as e: + raise ValueError(f"Failed to download or decode image from URL: {url!r}") 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`.""" @@ -1498,9 +1583,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 cb0be5823..57cb2207c 100644 --- a/mellea/stdlib/components/chat.py +++ b/mellea/stdlib/components/chat.py @@ -38,8 +38,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). documents (list[Document] | None): Optional documents associated with the message. diff --git a/test/backends/test_vision_ollama.py b/test/backends/test_vision_ollama.py index 9b5288e8e..173b0e023 100644 --- a/test/backends/test_vision_ollama.py +++ b/test/backends/test_vision_ollama.py @@ -132,12 +132,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.backends.ollama._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.backends.ollama._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_image_block_in_chat(mocked_session: MelleaSession, pil_image: Image.Image): image_block = ImageBlock.from_pil_image(pil_image) diff --git a/test/core/test_base.py b/test/core/test_base.py index 78055011f..7b30228ac 100644 --- a/test/core/test_base.py +++ b/test/core/test_base.py @@ -15,6 +15,7 @@ ModelOutputThunk, RawProviderResponse, blockify, + make_image_block, ) from mellea.core.backend import generate_walk from mellea.stdlib.components import Message @@ -173,6 +174,92 @@ def test_message_mixed_image_types(): assert len(msg.images) == 2 # type: ignore[arg-type] +# --- 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"} + + +def test_make_image_block_url_convert_to_base64(monkeypatch): + from contextlib import contextmanager + + @contextmanager + def fake_urlopen(url): + class _Resp: + def read(self_inner): + return _png_bytes() + + yield _Resp() + + import mellea.core.base as base_mod + + monkeypatch.setattr(base_mod.urllib.request, "urlopen", fake_urlopen) + + 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 urllib.error + + def boom(url): + raise urllib.error.URLError("no network") + + import mellea.core.base as base_mod + + monkeypatch.setattr(base_mod.urllib.request, "urlopen", 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_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 --- From 7543c4e34704d8cce6260c5b075638dd2ab9bc43 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Mon, 13 Jul 2026 13:50:56 -0400 Subject: [PATCH 2/3] review feedback Signed-off-by: AngeloDanducci --- mellea/backends/ollama.py | 25 ++++-- mellea/core/base.py | 86 +++++++++++++++++-- test/backends/test_vision_ollama.py | 53 ++++++++++++ test/core/test_base.py | 124 ++++++++++++++++++++++++---- 4 files changed, 258 insertions(+), 30 deletions(-) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index 3227d6a21..b0b60e2ef 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -386,11 +386,13 @@ 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` whose image - cannot be downloaded or decoded; Ollama requires base64-encoded - images, so URL images are fetched and encoded automatically. + cannot be downloaded or decoded. """ # Start by awaiting any necessary computation. await self.do_generate_walk(action) @@ -425,13 +427,20 @@ async def generate_from_chat_context( image_values: list[str] | None = None if m.images is not None: # Ollama only accepts base64-encoded images, so URL images are - # downloaded and encoded on the fly rather than rejected. - image_values = [ - _download_image_as_base64(str(img.value)) + # downloaded and encoded on the fly rather than rejected. 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(_download_image_as_base64, str(img.value)) + for i, img in enumerate(m.images) if isinstance(img, ImageUrlBlock) - else str(img.value) - for img in m.images - ] + } + if url_downloads: + downloaded = await asyncio.gather(*url_downloads.values()) + for i, value in zip(url_downloads.keys(), downloaded): + image_values[i] = value conversation.append( { "role": m.role, diff --git a/mellea/core/base.py b/mellea/core/base.py index 514ef221b..c47b012dc 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -19,9 +19,10 @@ import binascii import datetime import enum +import ipaddress import logging -import urllib.error -import urllib.request +import socket +import urllib.parse from collections.abc import Callable, Coroutine, Iterable, Mapping from copy import copy, deepcopy from dataclasses import dataclass, field @@ -40,6 +41,7 @@ if TYPE_CHECKING: import torch +import requests import typing_extensions from PIL import Image as PILImage @@ -222,6 +224,49 @@ def __repr__(self) -> str: return f"ImageUrlBlock({self.value}, {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.""" + + +def _assert_public_url(url: str) -> None: + """Reject a URL whose host resolves to a non-public IP address. + + Resolves every address the host maps to and rejects the download if any is + private, loopback, link-local, reserved, or multicast. This blocks + server-side request forgery (SSRF) against cloud metadata endpoints + (`169.254.169.254`), localhost services, and RFC 1918 ranges. + + Args: + url: The `http://`/`https://` URL to validate. + + Raises: + ValueError: If the host is missing or resolves to a non-public IP. + """ + host = urllib.parse.urlsplit(url).hostname + if not host: + raise ValueError(f"URL has no host to validate: {url!r}") + try: + infos = socket.getaddrinfo(host, None) + except socket.gaierror as e: + raise ValueError(f"Could not resolve host for URL: {url!r}") from e + for info in infos: + ip = ipaddress.ip_address(info[4][0]) + if ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_reserved + or ip.is_multicast + or ip.is_unspecified + ): + raise ValueError( + f"Refusing to download from non-public address {ip} for URL: {url!r}" + ) + + def _download_image_as_base64(url: str) -> str: """Download an image from a URL and return it as a base64-encoded PNG string. @@ -229,6 +274,12 @@ def _download_image_as_base64(url: str) -> str: real image, and re-encodes the result as a base64 PNG so the output is consistent with `ImageBlock`'s expected format. + The download is hardened against abuse: the host is validated against + non-public IP ranges before connecting (SSRF), redirects are refused so the + validation cannot be bypassed, 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. @@ -236,15 +287,34 @@ def _download_image_as_base64(url: str) -> str: str: The base64-encoded PNG representation of the downloaded image. Raises: - ValueError: If the image cannot be downloaded or the downloaded bytes - cannot be decoded as an image. + ValueError: If the host resolves to a non-public address, the response + exceeds the size cap, or the image cannot be downloaded or decoded. """ + _assert_public_url(url) try: - with urllib.request.urlopen(url) as response: # scheme validated by caller - raw = response.read() + with requests.get( # scheme validated by caller + url, + timeout=_IMAGE_DOWNLOAD_TIMEOUT_S, + allow_redirects=False, # a redirect could bypass the IP guard + 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 (urllib.error.URLError, OSError, ValueError) as e: - raise ValueError(f"Failed to download or decode image from URL: {url!r}") from e + 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) diff --git a/test/backends/test_vision_ollama.py b/test/backends/test_vision_ollama.py index 173b0e023..853ea610f 100644 --- a/test/backends/test_vision_ollama.py +++ b/test/backends/test_vision_ollama.py @@ -172,6 +172,59 @@ def test_image_url_block_download_failure_raises(mocked_session: MelleaSession): mocked_session.chat("What is in this image?", images=images) +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` + `getaddrinfo`) so the + import alias, 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.socket.getaddrinfo", + return_value=[(2, 1, 6, "", ("93.184.216.34", 0))], + ), + 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 7b30228ac..3ad1edc72 100644 --- a/test/core/test_base.py +++ b/test/core/test_base.py @@ -216,40 +216,136 @@ def test_make_image_block_preserves_meta(): assert block._meta == {"alt": "a dot"} -def test_make_image_block_url_convert_to_base64(monkeypatch): - from contextlib import contextmanager +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.""" - @contextmanager - def fake_urlopen(url): - class _Resp: - def read(self_inner): - return _png_bytes() + 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} + ) - yield _Resp() + def raise_for_status(self): + return None + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _patch_download(monkeypatch, fake_get): + """Bypass the SSRF guard and stub `requests.get` for a download test.""" import mellea.core.base as base_mod - monkeypatch.setattr(base_mod.urllib.request, "urlopen", fake_urlopen) + monkeypatch.setattr(base_mod, "_assert_public_url", lambda url: None) + 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 urllib.error +def test_make_image_block_download_disables_redirects(monkeypatch): + """The download must pass allow_redirects=False so a 302 can't bypass the guard.""" + seen: dict = {} + + def fake_get(url, **kw): + seen.update(kw) + return _FakeResponse(_png_bytes()) - def boom(url): - raise urllib.error.URLError("no network") + _patch_download(monkeypatch, fake_get) + make_image_block("https://example.com/cat.png", convert_to_base64=True) + assert seen.get("allow_redirects") is False + +def test_make_image_block_download_failure_raises(monkeypatch): import mellea.core.base as base_mod - monkeypatch.setattr(base_mod.urllib.request, "urlopen", boom) + 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_assert_public_url_rejects_private_ip(monkeypatch): + from mellea.core import base as base_mod + + # getaddrinfo returns (family, type, proto, canonname, sockaddr); sockaddr[0] is the IP. + monkeypatch.setattr( + base_mod.socket, + "getaddrinfo", + lambda host, port: [(2, 1, 6, "", ("169.254.169.254", 0))], + ) + with pytest.raises(ValueError, match="non-public"): + base_mod._assert_public_url("http://metadata.example/latest") + + +def test_assert_public_url_rejects_loopback(monkeypatch): + from mellea.core import base as base_mod + + monkeypatch.setattr( + base_mod.socket, + "getaddrinfo", + lambda host, port: [(2, 1, 6, "", ("127.0.0.1", 0))], + ) + with pytest.raises(ValueError, match="non-public"): + base_mod._assert_public_url("http://localhost/x.png") + + +def test_assert_public_url_allows_public_ip(monkeypatch): + from mellea.core import base as base_mod + + monkeypatch.setattr( + base_mod.socket, + "getaddrinfo", + lambda host, port: [(2, 1, 6, "", ("93.184.216.34", 0))], + ) + # Should not raise. + base_mod._assert_public_url("https://example.com/cat.png") + + def test_make_image_block_invalid_string_raises(): with pytest.raises(ValueError, match="could not interpret"): make_image_block("not-a-url-or-base64!!!") From bdb00ee3163098a9c14aa5cbf0133a0bcf62bb34 Mon Sep 17 00:00:00 2001 From: AngeloDanducci Date: Wed, 15 Jul 2026 14:32:25 -0400 Subject: [PATCH 3/3] remove ssrf guard add url keyed imageurlblock caching Signed-off-by: AngeloDanducci --- mellea/backends/ollama.py | 14 ++-- mellea/core/base.py | 108 +++++++++++++++++----------- test/backends/test_vision_ollama.py | 28 ++++---- test/core/test_base.py | 84 +++++++++++----------- 4 files changed, 131 insertions(+), 103 deletions(-) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index b0b60e2ef..167cff54f 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -25,7 +25,7 @@ ModelToolCall, RawProviderResponse, ) -from ..core.base import AbstractMelleaTool, _download_image_as_base64 +from ..core.base import AbstractMelleaTool from ..formatters import ChatFormatter, TemplateFormatter from ..helpers import ( DEFAULT_CHUNK_TIMEOUT, @@ -427,13 +427,15 @@ async def generate_from_chat_context( image_values: list[str] | None = None if m.images is not None: # Ollama only accepts base64-encoded images, so URL images are - # downloaded and encoded on the fly rather than rejected. 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. + # 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(_download_image_as_base64, str(img.value)) + i: asyncio.to_thread(img.resolve_base64) for i, img in enumerate(m.images) if isinstance(img, ImageUrlBlock) } diff --git a/mellea/core/base.py b/mellea/core/base.py index c47b012dc..b00e77dd5 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -19,10 +19,9 @@ import binascii import datetime import enum -import ipaddress import logging -import socket -import urllib.parse +import threading +from collections import OrderedDict from collections.abc import Callable, Coroutine, Iterable, Mapping from copy import copy, deepcopy from dataclasses import dataclass, field @@ -219,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__()})" @@ -230,41 +248,51 @@ def __repr__(self) -> str: _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 _assert_public_url(url: str) -> None: - """Reject a URL whose host resolves to a non-public IP address. +def _cached_download_image_as_base64(url: str) -> str: + """Download an image as base64, memoizing the result per URL. - Resolves every address the host maps to and rejects the download if any is - private, loopback, link-local, reserved, or multicast. This blocks - server-side request forgery (SSRF) against cloud metadata endpoints - (`169.254.169.254`), localhost services, and RFC 1918 ranges. + 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: The `http://`/`https://` URL to validate. + 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 host is missing or resolves to a non-public IP. + ValueError: If the response exceeds the size cap or the image cannot + be downloaded or decoded. """ - host = urllib.parse.urlsplit(url).hostname - if not host: - raise ValueError(f"URL has no host to validate: {url!r}") - try: - infos = socket.getaddrinfo(host, None) - except socket.gaierror as e: - raise ValueError(f"Could not resolve host for URL: {url!r}") from e - for info in infos: - ip = ipaddress.ip_address(info[4][0]) - if ( - ip.is_private - or ip.is_loopback - or ip.is_link_local - or ip.is_reserved - or ip.is_multicast - or ip.is_unspecified - ): - raise ValueError( - f"Refusing to download from non-public address {ip} for URL: {url!r}" - ) + 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: @@ -274,11 +302,9 @@ def _download_image_as_base64(url: str) -> str: real image, and re-encodes the result as a base64 PNG so the output is consistent with `ImageBlock`'s expected format. - The download is hardened against abuse: the host is validated against - non-public IP ranges before connecting (SSRF), redirects are refused so the - validation cannot be bypassed, 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`. + 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. @@ -287,16 +313,12 @@ def _download_image_as_base64(url: str) -> str: str: The base64-encoded PNG representation of the downloaded image. Raises: - ValueError: If the host resolves to a non-public address, the response - exceeds the size cap, or the image cannot be downloaded or decoded. + ValueError: If the response exceeds the size cap or the image cannot + be downloaded or decoded. """ - _assert_public_url(url) try: with requests.get( # scheme validated by caller - url, - timeout=_IMAGE_DOWNLOAD_TIMEOUT_S, - allow_redirects=False, # a redirect could bypass the IP guard - stream=True, + url, timeout=_IMAGE_DOWNLOAD_TIMEOUT_S, stream=True ) as response: response.raise_for_status() declared = response.headers.get("Content-Length") diff --git a/test/backends/test_vision_ollama.py b/test/backends/test_vision_ollama.py index 853ea610f..3945bf455 100644 --- a/test/backends/test_vision_ollama.py +++ b/test/backends/test_vision_ollama.py @@ -39,6 +39,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) @@ -142,7 +152,7 @@ def test_image_url_block_auto_downloaded_by_ollama( images: list[ImageBlock | ImageUrlBlock] = [url_block] with patch( - "mellea.backends.ollama._download_image_as_base64", return_value=encoded + "mellea.core.base._download_image_as_base64", return_value=encoded ) as mock_download: mocked_session.chat("What is in this image?", images=images) @@ -165,7 +175,7 @@ def test_image_url_block_download_failure_raises(mocked_session: MelleaSession): images: list[ImageBlock | ImageUrlBlock] = [url_block] with patch( - "mellea.backends.ollama._download_image_as_base64", + "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"): @@ -177,9 +187,9 @@ def test_image_url_block_drives_real_download( ): """Exercise the real `_download_image_as_base64` through the Ollama path. - Patches only the network layer (`requests.get` + `getaddrinfo`) so the - import alias, thread-offload, and payload wiring are all covered — this - catches regressions the fully-mocked download tests cannot. + 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") @@ -205,13 +215,7 @@ def __enter__(self): def __exit__(self, *exc): return False - with ( - patch( - "mellea.core.base.socket.getaddrinfo", - return_value=[(2, 1, 6, "", ("93.184.216.34", 0))], - ), - patch("mellea.core.base.requests.get", return_value=_FakeResponse()), - ): + 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() diff --git a/test/core/test_base.py b/test/core/test_base.py index 3ad1edc72..0e2cf946d 100644 --- a/test/core/test_base.py +++ b/test/core/test_base.py @@ -245,11 +245,20 @@ 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): - """Bypass the SSRF guard and stub `requests.get` for a download test.""" + """Stub `requests.get` for a download test.""" import mellea.core.base as base_mod - monkeypatch.setattr(base_mod, "_assert_public_url", lambda url: None) monkeypatch.setattr(base_mod.requests, "get", fake_get) @@ -261,19 +270,6 @@ def test_make_image_block_url_convert_to_base64(monkeypatch): assert ImageBlock.is_valid_base64_png(str(block)) -def test_make_image_block_download_disables_redirects(monkeypatch): - """The download must pass allow_redirects=False so a 302 can't bypass the guard.""" - seen: dict = {} - - def fake_get(url, **kw): - seen.update(kw) - return _FakeResponse(_png_bytes()) - - _patch_download(monkeypatch, fake_get) - make_image_block("https://example.com/cat.png", convert_to_base64=True) - assert seen.get("allow_redirects") is False - - def test_make_image_block_download_failure_raises(monkeypatch): import mellea.core.base as base_mod @@ -309,41 +305,45 @@ def test_make_image_block_download_rejects_oversized_content_length(monkeypatch) make_image_block("https://example.com/big.png", convert_to_base64=True) -def test_assert_public_url_rejects_private_ip(monkeypatch): - from mellea.core import base as base_mod +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] = {} - # getaddrinfo returns (family, type, proto, canonname, sockaddr); sockaddr[0] is the IP. - monkeypatch.setattr( - base_mod.socket, - "getaddrinfo", - lambda host, port: [(2, 1, 6, "", ("169.254.169.254", 0))], - ) - with pytest.raises(ValueError, match="non-public"): - base_mod._assert_public_url("http://metadata.example/latest") + def fake_get(url, **kw): + calls[url] = calls.get(url, 0) + 1 + return _FakeResponse(_png_bytes()) + _patch_download(monkeypatch, fake_get) -def test_assert_public_url_rejects_loopback(monkeypatch): - from mellea.core import base as base_mod + 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() - monkeypatch.setattr( - base_mod.socket, - "getaddrinfo", - lambda host, port: [(2, 1, 6, "", ("127.0.0.1", 0))], - ) - with pytest.raises(ValueError, match="non-public"): - base_mod._assert_public_url("http://localhost/x.png") + 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_assert_public_url_allows_public_ip(monkeypatch): + +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 - monkeypatch.setattr( - base_mod.socket, - "getaddrinfo", - lambda host, port: [(2, 1, 6, "", ("93.184.216.34", 0))], - ) - # Should not raise. - base_mod._assert_public_url("https://example.com/cat.png") + _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():