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..b404bf2701c 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,98 @@ 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=", + ], +) +def test_cmc_rejects_invalid_image_data_uri( + uri: str, + ollama_unit_test_env: dict[str, str], +) -> None: + message = Message( + contents=[Content.from_uri(uri=uri, media_type="image/png")], + role="user", + ) + ollama_client = OllamaChatClient() + + with pytest.raises(ChatClientInvalidRequestException): + ollama_client._format_user_message(message) + + +@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)