From 99ac95952191d4e68940daecee264a56c2ce81e3 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 16 Oct 2026 17:13:33 +0200 Subject: [PATCH 1/2] Python: decode Ollama inline images before serialization --- .../agent_framework_ollama/_chat_client.py | 29 +++++- .../ollama/tests/test_ollama_chat_client.py | 97 ++++++++++++++++++- 2 files changed, 122 insertions(+), 4 deletions(-) diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index a0bcc046296..970840ed6e4 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -2,6 +2,8 @@ from __future__ import annotations +import base64 +import binascii import json import logging import sys @@ -45,6 +47,7 @@ # Rename imported types to avoid naming conflicts with Agent Framework types from ollama._types import ChatResponse as OllamaChatResponse +from ollama._types import Image as OllamaImage from ollama._types import Message as OllamaMessage from pydantic import BaseModel @@ -533,10 +536,32 @@ def _format_user_message(self, message: Message) -> list[OllamaMessage]: raise ChatClientInvalidRequestException( "Only image data content is supported for user messages in Ollama." ) - # Ollama expects base64 strings without prefix - user_message["images"] = [c.uri.split(",")[1] for c in data_contents if c.uri] + user_message["images"] = [OllamaImage(value=self._decode_image_data_uri(c)) for c in data_contents] return [user_message] + @staticmethod + def _decode_image_data_uri(content: Content) -> bytes: + uri = content.uri + if not uri or not uri.startswith("data:") or "," not in uri: + raise ChatClientInvalidRequestException("Ollama image content must be a base64-encoded data:image URI.") + + metadata, payload = uri.split(",", 1) + metadata_parts = metadata[5:].split(";") + parameters = metadata_parts[1:] + if ( + not metadata_parts[0].lower().startswith("image/") + or not parameters + or parameters[-1] != "base64" + or any(not parameter or "=" not in parameter for parameter in parameters[:-1]) + or not payload + ): + raise ChatClientInvalidRequestException("Ollama image content must be a base64-encoded data:image URI.") + + try: + return base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError) as exc: + raise ChatClientInvalidRequestException("Ollama image content contains invalid base64 data.") from exc + def _format_assistant_message(self, message: Message) -> list[OllamaMessage]: text_content = message.text # Ollama shouldn't have encrypted reasoning, so we just process text. diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index ab185cb3ce1..ea8307d1289 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -1,8 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. +import base64 import os import uuid from collections.abc import AsyncIterable +from pathlib import Path from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -25,6 +27,7 @@ ) from ollama import AsyncClient from ollama._types import ChatResponse as OllamaChatResponse +from ollama._types import Image as OllamaImage from ollama._types import Message as OllamaMessage from openai import AsyncStream from pydantic import BaseModel @@ -634,18 +637,108 @@ async def test_cmc_with_data_content_type( chat_history: list[Message], mock_chat_completion_response: OllamaChatResponse, ) -> None: + image_bytes = b"\x89PNG\r\n\x1a\n" mock_chat.return_value = mock_chat_completion_response chat_history.append( Message( - contents=[Content.from_uri(uri="data:image/png;base64,xyz", media_type="image/png")], + contents=[ + Content.from_uri( + uri=f"data:image/png;base64,{base64.b64encode(image_bytes).decode()}", + media_type="image/png", + ) + ], role="user", ) ) ollama_client = OllamaChatClient() - result = await ollama_client.get_response(messages=chat_history) + with ( + patch.object(Path, "exists", side_effect=AssertionError("Path.exists must not be called")) as path_exists, + patch.object( + Path, "read_bytes", side_effect=AssertionError("Path.read_bytes must not be called") + ) as read_bytes, + ): + result = await ollama_client.get_response(messages=chat_history) + sent_image = mock_chat.call_args.kwargs["messages"][0].images[0] + assert isinstance(sent_image, OllamaImage) + assert sent_image.value == image_bytes + assert sent_image.model_dump() == base64.b64encode(image_bytes).decode() + assert result.text == "test" + path_exists.assert_not_called() + read_bytes.assert_not_called() + + +@pytest.mark.parametrize( + "uri", + [ + "data:image/png,/tmp/image.png", + "data:image/png;base64,/tmp/image.png", + r"data:image/png;base64,C:\temp\image.png", + r"data:image/png;base64,\\server\share\image.png", + "data:image/png;base64,../image.png", + "data:image/png;base64,", + "data:text/plain;base64,aGVsbG8=", + ], +) +@patch.object(AsyncClient, "chat", new_callable=AsyncMock) +async def test_cmc_rejects_invalid_image_data_uri( + mock_chat: AsyncMock, + uri: str, + ollama_unit_test_env: dict[str, str], + chat_history: list[Message], + mock_chat_completion_response: OllamaChatResponse, +) -> None: + mock_chat.return_value = mock_chat_completion_response + chat_history.append( + Message( + contents=[Content.from_uri(uri=uri, media_type="image/png")], + role="user", + ) + ) + + ollama_client = OllamaChatClient() + + with pytest.raises(ChatClientInvalidRequestException): + await ollama_client.get_response(messages=chat_history) + + mock_chat.assert_not_called() + + +@patch.object(AsyncClient, "chat", new_callable=AsyncMock) +async def test_cmc_treats_base64_shaped_path_as_inline_bytes( + mock_chat: AsyncMock, + ollama_unit_test_env: dict[str, str], + chat_history: list[Message], + mock_chat_completion_response: OllamaChatResponse, +) -> None: + payload = "////tmpx" + expected_bytes = base64.b64decode(payload, validate=True) + mock_chat.return_value = mock_chat_completion_response + chat_history.append( + Message( + contents=[Content.from_uri(uri=f"data:image/png;base64,{payload}", media_type="image/png")], + role="user", + ) + ) + + ollama_client = OllamaChatClient() + + with ( + patch.object(Path, "exists", side_effect=AssertionError("Path.exists must not be called")) as path_exists, + patch.object( + Path, "read_bytes", side_effect=AssertionError("Path.read_bytes must not be called") + ) as read_bytes, + ): + await ollama_client.get_response(messages=chat_history) + sent_image = mock_chat.call_args.kwargs["messages"][0].images[0] + assert isinstance(sent_image, OllamaImage) + assert sent_image.value == expected_bytes + assert sent_image.model_dump() == payload + + path_exists.assert_not_called() + read_bytes.assert_not_called() @patch.object(AsyncClient, "chat", new_callable=AsyncMock) From 3386ea0e12147ed7156ea1295d1849985bbbed36 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 16 Oct 2026 17:23:29 +0200 Subject: [PATCH 2/2] Test Ollama URI validation at connector boundary --- .../ollama/tests/test_ollama_chat_client.py | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index ea8307d1289..b404bf2701c 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -682,28 +682,18 @@ async def test_cmc_with_data_content_type( "data:text/plain;base64,aGVsbG8=", ], ) -@patch.object(AsyncClient, "chat", new_callable=AsyncMock) -async def test_cmc_rejects_invalid_image_data_uri( - mock_chat: AsyncMock, +def test_cmc_rejects_invalid_image_data_uri( uri: str, ollama_unit_test_env: dict[str, str], - chat_history: list[Message], - mock_chat_completion_response: OllamaChatResponse, ) -> None: - mock_chat.return_value = mock_chat_completion_response - chat_history.append( - Message( - contents=[Content.from_uri(uri=uri, media_type="image/png")], - role="user", - ) + message = Message( + contents=[Content.from_uri(uri=uri, media_type="image/png")], + role="user", ) - ollama_client = OllamaChatClient() with pytest.raises(ChatClientInvalidRequestException): - await ollama_client.get_response(messages=chat_history) - - mock_chat.assert_not_called() + ollama_client._format_user_message(message) @patch.object(AsyncClient, "chat", new_callable=AsyncMock)