Python: Decode Ollama inline image data before serialization

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.

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)