feat: imageblock and imgurblock ergonomics#1377
Conversation
| return ImageBlock.pil_to_base64(image) | ||
|
|
||
|
|
||
| def make_image_block( |
There was a problem hiding this comment.
is it worth debug logging when a conversion happens?
psschwei
left a comment
psschwei
left a comment
There was a problem hiding this comment.
The make_image_block factory is clean, well-tested, and mergeable as-is. But the Ollama auto-download change swaps an explicit opt-in ValueError for an automatic, synchronous network fetch inside an async method — with no timeout, no size cap, and no SSRF protection. That needs to be addressed (or explicitly accepted by a maintainer) before merge.
Verified locally: tests pass (16/16 in the touched suites), ruff format/ruff check clean, and the docstring quality gate passes at 100%.
🔴 Blocking
1. Blocking urlopen inside an async coroutine — stalls the event loop (mellea/core/base.py:243, called from mellea/backends/ollama.py:429)
_download_image_as_base64 does a fully synchronous urllib.request.urlopen(...).read() and is called directly inside async def generate_from_chat_context — no asyncio.to_thread, no executor. One chat request with a URL image blocks the entire event loop for the whole download, freezing every other concurrent generation sharing that loop (m serve, multi-session apps). Offload it: await asyncio.to_thread(_download_image_as_base64, url).
2. No timeout on urlopen (mellea/core/base.py:243)
urlopen(url) has no timeout=. A slow/hung remote can block indefinitely — and per #1, it blocks the shared loop while doing so. This is now reachable by default (it replaced the old ValueError). Pass an explicit timeout= and fold TimeoutError/socket.timeout into the existing ValueError wrapper.
3. No response size cap (mellea/core/base.py:243-247)
response.read() buffers the whole body into memory before PIL ever sees it. A crafted/infinite response is a memory-exhaustion DoS, triggered automatically by the Ollama path. Cap it (check Content-Length, or read(MAX + 1) and reject on overflow).
4. SSRF exposure — now automatic, not opt-in (mellea/core/base.py:225-248)
Before this PR, nothing in mellea made an outbound GET from a URL embedded in message content. Now any ImageUrlBlock routed through Ollama causes the process to fetch that URL — including URLs from untrusted sources (documents, tool results, model output echoed back). No scheme/host/IP restrictions (169.254.169.254 metadata, localhost:6379, RFC1918), and urlopen follows redirects, so an input-string host check would be bypassable via a 302. At minimum this warrants explicit maintainer/security sign-off; a reasonable mitigation is validating resolved IPs against private/loopback/link-local ranges and disabling redirects.
🟡 Suggestions (non-blocking)
- Error message loses the cause (
base.py:246) —from epreserves the chain, but the surfaced message omits why (404 vs DNS vs bad image). Include{e}. - Sequential downloads per message (
ollama.py:429-434) — NImageUrlBlocks = N sequential blocking round-trips. Once #1 is async,asyncio.gatherthem (keepreturn_exceptions=Falseper AGENTS.md §9). Raises:docstring is padded with prose (ollama.py:390-393) — the "fetched and encoded automatically" sentence belongs in the body, not theRaises:entry.- Naming — issue #1312 suggested
image_block; PR shipsmake_image_block. Fine (repo hasmake_execution_environmentprecedent), but worth a one-line note in the PR so it doesn't read as a misread spec. - Ollama failure test is fully mocked (
test_vision_ollama.py:163) — no test drives the real_download_image_as_base64through the Ollama path. Consider one that patches onlyurllib.request.urlopen(astest_base.py:219already does) to catch import-alias regressions.
✅ Done well
make_image_blocktests cover all four dispatch branches, both error paths, and meta propagation — thorough.- Ollama change is a minimal diff that reuses
_strip_data_uri_prefixand the existing payload shape. - Docstring double→single-backtick fixes correctly follow AGENTS.md §5 and exactly match what #1312 asked for.
- The
"passed directly to OpenAI-compatible backends"docstring claim is accurate — verified againstopenai_compatible_helpers.py:200-206, which passesImageUrlBlock.valuestraight through asimage_url.url. - Quality gate (§10.6) passes at 100% for the new public API, including
Raises:/Returns:.
Items #1–#4 are all in the ~20-line _download_image_as_base64 helper, so a single focused revision (thread offload + timeout + size cap + IP/redirect guard) clears the whole blocking set.
|
Should be ready for re-review, I think the only thing not changed was regarding naming. I chose |
psschwei
left a comment
psschwei
left a comment
There was a problem hiding this comment.
LGTM
In the interest of not accidentally busting core (albeit probably overly cautious), would be good to get an additional set of eyes on this before merging. This basically mirrors the existing HF / OpenAI process, so shouldn't be a problem, but "with great power..."
| 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}" | ||
| ) |
There was a problem hiding this comment.
I'm not sure I agree with this. Similar to the discussion on logging, I don't know if we are the correct place to be blocking this. I could image you have a local server or something that you want to pull images from (or test pulling images from).
There was a problem hiding this comment.
It also looks like including this doesn't necessarily fix the issue since the DNS resolution could differ between this point and the actual requests.get fetch. A motivated attacker could change these between calls unless we pin the IP.
There was a problem hiding this comment.
Removed the ssrf guard here, agreed with local image host. Since this isn't the first time we've seen something regarding ssrf pop up do you think that's something we want to add or should it be on the user to have network level egress filter?
There was a problem hiding this comment.
I'm kinda torn on this. I feel like we should prevent users from doing dangerous things but not limit them too much. I'm not sure what the heuristic for that is / what the actual rules for that are. I think here, it's fine to remove the ssrf guard.
| 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. | ||
|
|
||
| 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. | ||
|
|
||
| Returns: | ||
| 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. | ||
| """ | ||
| _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, | ||
| ) 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) |
There was a problem hiding this comment.
Won't this function get called for each conversation turn for each ImageUrlBlock in that conversation (for Ollama at least?)? I think we should add a field to ImageUrlBlock to indicate if it's already been downloaded / store the base64 version / etc... so that we effectively cache this work.
There was a problem hiding this comment.
Added a url keyd cache to avoid this.
|
Rebased to fix conflicts. |
Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>
Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>
Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>
|
Aaaand another rebase after the thinking merge. |
a846b04
* feat: imageblock and imgurblock ergonomics Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com> * review feedback Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com> * remove ssrf guard add url keyed imageurlblock caching Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com> --------- Signed-off-by: AngeloDanducci <angelo.danducci.ii@ibm.com>
3 participants
Pull Request
Issue
Fixes #1312
Description
Adds ImageBlock and ImageUrBlock ergonomics.
Note: I chose
make_image_blocktoavoid image_blockandImageBlockbeing a bit too close when reading quickly.Testing
Attribution
Adding a new component, requirement, sampling strategy, or tool?
If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.
NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.