diff --git a/python/semantic_kernel/connectors/ai/README.md b/python/semantic_kernel/connectors/ai/README.md index 6fe7510a0697..352aed81634d 100644 --- a/python/semantic_kernel/connectors/ai/README.md +++ b/python/semantic_kernel/connectors/ai/README.md @@ -48,5 +48,6 @@ All base clients inherit from the [`AIServiceClientBase`](../../services/ai_serv | Ollama | [`OllamaChatCompletion`](./ollama/services/ollama_chat_completion.py) | | | [`OllamaTextCompletion`](./ollama/services/ollama_text_completion.py) | | | [`OllamaTextEmbedding`](./ollama/services/ollama_text_embedding.py) | +| | [`OllamaTextToImage`](./ollama/services/ollama_text_to_image.py) | | Onnx | [`OnnxGenAIChatCompletion`](./onnx/services/onnx_gen_ai_chat_completion.py) | | | [`OnnxGenAITextCompletion`](./onnx/services/onnx_gen_ai_text_completion.py) | \ No newline at end of file diff --git a/python/semantic_kernel/connectors/ai/ollama/__init__.py b/python/semantic_kernel/connectors/ai/ollama/__init__.py index 1150a5e46f76..5aba3c51102f 100644 --- a/python/semantic_kernel/connectors/ai/ollama/__init__.py +++ b/python/semantic_kernel/connectors/ai/ollama/__init__.py @@ -5,10 +5,12 @@ OllamaEmbeddingPromptExecutionSettings, OllamaPromptExecutionSettings, OllamaTextPromptExecutionSettings, + OllamaTextToImagePromptExecutionSettings, ) from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion from semantic_kernel.connectors.ai.ollama.services.ollama_text_completion import OllamaTextCompletion from semantic_kernel.connectors.ai.ollama.services.ollama_text_embedding import OllamaTextEmbedding +from semantic_kernel.connectors.ai.ollama.services.ollama_text_to_image import OllamaTextToImage __all__ = [ "OllamaChatCompletion", @@ -18,4 +20,6 @@ "OllamaTextCompletion", "OllamaTextEmbedding", "OllamaTextPromptExecutionSettings", + "OllamaTextToImage", + "OllamaTextToImagePromptExecutionSettings", ] diff --git a/python/semantic_kernel/connectors/ai/ollama/ollama_prompt_execution_settings.py b/python/semantic_kernel/connectors/ai/ollama/ollama_prompt_execution_settings.py index 6139356e2a28..963adf28b423 100644 --- a/python/semantic_kernel/connectors/ai/ollama/ollama_prompt_execution_settings.py +++ b/python/semantic_kernel/connectors/ai/ollama/ollama_prompt_execution_settings.py @@ -37,3 +37,11 @@ class OllamaChatPromptExecutionSettings(OllamaPromptExecutionSettings): class OllamaEmbeddingPromptExecutionSettings(OllamaPromptExecutionSettings): """Settings for Ollama embedding prompt execution.""" + + +class OllamaTextToImagePromptExecutionSettings(OllamaPromptExecutionSettings): + """Settings for Ollama text to image execution.""" + + width: int | None = None + height: int | None = None + steps: int | None = None diff --git a/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py b/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py index ee4fa05863a4..c51545e1ac13 100644 --- a/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py +++ b/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py @@ -20,6 +20,7 @@ class OllamaSettings(KernelBaseSettings): - chat_model_id: str - The chat model ID. (Env var OLLAMA_CHAT_MODEL_ID) - text_model_id: str - The text model ID. (Env var OLLAMA_TEXT_MODEL_ID) - embedding_model_id: str - The embedding model ID. (Env var OLLAMA_EMBEDDING_MODEL_ID) + - image_model_id: str - The image generation model ID. (Env var OLLAMA_IMAGE_MODEL_ID) Optional settings for prefix 'OLLAMA' are: - host: HttpsUrl - The endpoint of the Ollama service. (Env var OLLAMA_HOST) @@ -30,4 +31,5 @@ class OllamaSettings(KernelBaseSettings): chat_model_id: str | None = None text_model_id: str | None = None embedding_model_id: str | None = None + image_model_id: str | None = None host: str | None = None diff --git a/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_to_image.py b/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_to_image.py new file mode 100644 index 000000000000..9ee4e206a0b8 --- /dev/null +++ b/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_to_image.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft. All rights reserved. + +import base64 +import logging +import sys +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any +from warnings import warn + +from ollama import AsyncClient +from pydantic import ValidationError + +from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import ( + OllamaTextToImagePromptExecutionSettings, +) +from semantic_kernel.connectors.ai.ollama.ollama_settings import OllamaSettings +from semantic_kernel.connectors.ai.ollama.services.ollama_base import OllamaBase +from semantic_kernel.connectors.ai.text_to_image_client_base import TextToImageClientBase +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError +from semantic_kernel.utils.feature_stage_decorator import experimental + +if TYPE_CHECKING: + from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings + +if sys.version_info >= (3, 12): + from typing import override # pragma: no cover +else: + from typing_extensions import override # pragma: no cover + +logger: logging.Logger = logging.getLogger(__name__) + + +@experimental +class OllamaTextToImage(OllamaBase, TextToImageClientBase): + """Ollama text to image client. + + Make sure to have the ollama service running either locally or remotely, with an + image generation model pulled, for example `x/z-image-turbo`. + """ + + def __init__( + self, + service_id: str | None = None, + ai_model_id: str | None = None, + host: str | None = None, + client: AsyncClient | None = None, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize an OllamaTextToImage service. + + Args: + service_id (Optional[str]): Service ID tied to the execution settings. (Optional) + ai_model_id (Optional[str]): The model name. (Optional) + host (Optional[str]): URL of the Ollama server, defaults to None and + will use the default Ollama service address: http://127.0.0.1:11434. (Optional) + client (Optional[AsyncClient]): A custom Ollama client to use for the service. (Optional) + env_file_path (str | None): Use the environment settings file as a fallback to using env vars. + env_file_encoding (str | None): The encoding of the environment settings file, defaults to 'utf-8'. + """ + try: + ollama_settings = OllamaSettings( + image_model_id=ai_model_id, + host=host, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + except ValidationError as ex: + raise ServiceInitializationError("Failed to create Ollama settings.", ex) from ex + + if not ollama_settings.image_model_id: + raise ServiceInitializationError("Ollama image model ID is not set.") + + super().__init__( + service_id=service_id or ollama_settings.image_model_id, + ai_model_id=ollama_settings.image_model_id, + client=client or AsyncClient(host=ollama_settings.host), + ) + + @override + async def generate_image( + self, + description: str, + width: int | None = None, + height: int | None = None, + settings: "PromptExecutionSettings | None" = None, + **kwargs: Any, + ) -> bytes: + """Generate an image from a text description. + + Args: + description: Description of the image. + width: Deprecated, use settings.width instead. + height: Deprecated, use settings.height instead. + settings: Execution settings for the prompt. + kwargs: Additional arguments passed to the Ollama generate endpoint. + + Returns: + bytes: The raw image bytes. + """ + image_settings = ( + OllamaTextToImagePromptExecutionSettings() + if settings is None + else OllamaTextToImagePromptExecutionSettings.from_prompt_execution_settings(settings) + ) + + if width is not None: + warn( + "The 'width' argument is deprecated. Use 'settings.width' instead.", + DeprecationWarning, + stacklevel=2, + ) + if image_settings.width is None: + image_settings.width = width + if height is not None: + warn( + "The 'height' argument is deprecated. Use 'settings.height' instead.", + DeprecationWarning, + stacklevel=2, + ) + if image_settings.height is None: + image_settings.height = height + + options = image_settings.prepare_settings_dict() + options.update(kwargs) + + response_object = await self.client.generate( + model=self.ai_model_id, + prompt=description, + stream=False, + **options, + ) + + image = getattr(response_object, "image", None) + if image is None and isinstance(response_object, Mapping): + image = response_object.get("image") + if not image: + raise ServiceInvalidResponseError( + "The Ollama response did not contain image data. Make sure the configured model " + f"('{self.ai_model_id}') is an image generation model." + ) + + return base64.b64decode(image) + + @override + def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]: + return OllamaTextToImagePromptExecutionSettings diff --git a/python/tests/unit/connectors/ai/ollama/conftest.py b/python/tests/unit/connectors/ai/ollama/conftest.py index 26155026db3c..d42adbf0827c 100644 --- a/python/tests/unit/connectors/ai/ollama/conftest.py +++ b/python/tests/unit/connectors/ai/ollama/conftest.py @@ -56,6 +56,7 @@ def ollama_unit_test_env(monkeypatch, host, exclude_list): "OLLAMA_CHAT_MODEL_ID": "test_chat_model_id", "OLLAMA_TEXT_MODEL_ID": "test_text_model_id", "OLLAMA_EMBEDDING_MODEL_ID": "test_embedding_model_id", + "OLLAMA_IMAGE_MODEL_ID": "test_image_model_id", "OLLAMA_HOST": host, } diff --git a/python/tests/unit/connectors/ai/ollama/services/test_ollama_text_to_image.py b/python/tests/unit/connectors/ai/ollama/services/test_ollama_text_to_image.py new file mode 100644 index 000000000000..e4608ef5c2d6 --- /dev/null +++ b/python/tests/unit/connectors/ai/ollama/services/test_ollama_text_to_image.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft. All rights reserved. + +import base64 +from unittest.mock import patch + +import pytest + +from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import ( + OllamaTextToImagePromptExecutionSettings, +) +from semantic_kernel.connectors.ai.ollama.services.ollama_text_to_image import OllamaTextToImage +from semantic_kernel.contents.image_content import ImageContent +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError + +ENCODED_IMAGE = base64.b64encode(b"test_image_bytes").decode() + + +def test_init_empty_service_id(model_id): + """Test that the service initializes correctly with an empty service id.""" + ollama = OllamaTextToImage(ai_model_id=model_id) + assert ollama.service_id == model_id + + +def test_custom_client(model_id, custom_client): + """Test that the service initializes correctly with a custom client.""" + ollama = OllamaTextToImage(ai_model_id=model_id, client=custom_client) + assert ollama.client == custom_client + + +def test_invalid_ollama_settings(): + """Test that the service initializes incorrectly with invalid settings.""" + with pytest.raises(ServiceInitializationError): + _ = OllamaTextToImage(ai_model_id=123) + + +@pytest.mark.parametrize("exclude_list", [["OLLAMA_IMAGE_MODEL_ID"]], indirect=True) +def test_init_empty_model_id(ollama_unit_test_env): + """Test that the service initializes incorrectly with an empty model id.""" + with pytest.raises(ServiceInitializationError): + _ = OllamaTextToImage(env_file_path="fake_env_file_path.env") + + +@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client +@patch("ollama.AsyncClient.generate") # mock_generate +async def test_custom_host(mock_generate, mock_client, model_id, host, prompt): + """Test that the service generates an image correctly with a custom host.""" + mock_generate.return_value = {"image": ENCODED_IMAGE} + + ollama = OllamaTextToImage(ai_model_id=model_id, host=host) + _ = await ollama.generate_image(prompt) + + mock_client.assert_called_once_with(host=host) + + +@patch("ollama.AsyncClient.generate") +async def test_generate_image(mock_generate, model_id, prompt): + """Test that the service decodes the base64 image returned by Ollama.""" + mock_generate.return_value = {"image": ENCODED_IMAGE} + settings = OllamaTextToImagePromptExecutionSettings() + settings.options = {"test_key": "test_value"} + + ollama = OllamaTextToImage(ai_model_id=model_id) + image = await ollama.generate_image(prompt, settings=settings) + + assert image == b"test_image_bytes" + mock_generate.assert_called_once_with( + model=model_id, + prompt=prompt, + stream=False, + options={"test_key": "test_value"}, + ) + + +@patch("ollama.AsyncClient.generate") +async def test_get_image_content(mock_generate, model_id, prompt): + """Test that the inherited get_image_content returns ImageContent with the image data.""" + mock_generate.return_value = {"image": ENCODED_IMAGE} + + ollama = OllamaTextToImage(ai_model_id=model_id) + content = await ollama.get_image_content(prompt, OllamaTextToImagePromptExecutionSettings()) + + assert isinstance(content, ImageContent) + assert content.data == b"test_image_bytes" + + +@patch("ollama.AsyncClient.generate") +async def test_generate_image_with_size_settings(mock_generate, model_id, prompt): + """Test that width, height and steps from the settings are forwarded to Ollama.""" + mock_generate.return_value = {"image": ENCODED_IMAGE} + settings = OllamaTextToImagePromptExecutionSettings(width=512, height=256, steps=4) + + ollama = OllamaTextToImage(ai_model_id=model_id) + _ = await ollama.generate_image(prompt, settings=settings) + + call_kwargs = mock_generate.call_args.kwargs + assert call_kwargs["width"] == 512 + assert call_kwargs["height"] == 256 + assert call_kwargs["steps"] == 4 + + +@patch("ollama.AsyncClient.generate") +async def test_generate_image_deprecated_width_and_height_arguments(mock_generate, model_id, prompt): + """Test that the deprecated width and height arguments still reach Ollama, with a warning.""" + mock_generate.return_value = {"image": ENCODED_IMAGE} + + ollama = OllamaTextToImage(ai_model_id=model_id) + with pytest.warns(DeprecationWarning): + image = await ollama.generate_image(prompt, width=512, height=256) + + assert image == b"test_image_bytes" + call_kwargs = mock_generate.call_args.kwargs + assert call_kwargs["width"] == 512 + assert call_kwargs["height"] == 256 + + +@patch("ollama.AsyncClient.generate") +async def test_generate_image_settings_take_precedence_over_arguments(mock_generate, model_id, prompt): + """Test that explicit settings win over the deprecated width and height arguments.""" + mock_generate.return_value = {"image": ENCODED_IMAGE} + settings = OllamaTextToImagePromptExecutionSettings(width=1024, height=1024) + + ollama = OllamaTextToImage(ai_model_id=model_id) + with pytest.warns(DeprecationWarning): + _ = await ollama.generate_image(prompt, width=512, height=256, settings=settings) + + call_kwargs = mock_generate.call_args.kwargs + assert call_kwargs["width"] == 1024 + assert call_kwargs["height"] == 1024 + + +@patch("ollama.AsyncClient.generate") +async def test_generate_image_without_image_in_response(mock_generate, model_id, prompt): + """Test that a response without image data raises instead of returning empty bytes.""" + mock_generate.return_value = {"response": "this model returns text"} + + ollama = OllamaTextToImage(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + await ollama.generate_image(prompt)