From 79207ffbbc6655391dd072b0afa994de407b6be0 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 16 Jul 2024 15:38:21 -0700 Subject: [PATCH] Python: Migrate to Ollama Python SDK (#7165) ### Motivation and Context ### Description Currently, the Ollama AI connector uses `aiohttp` sessions to interface with the local Ollama LLM service. However, we have made the decision to transition to the official Ollama Python SDK. Related to: https://github.com/microsoft/semantic-kernel/issues/6841, https://github.com/microsoft/semantic-kernel/issues/7134 Test coverage report: ![image](https://github.com/microsoft/semantic-kernel/assets/12570346/8d414f71-fc1b-4870-b038-0f7d2985b985) MyPy warnings: all cleared ### Contribution Checklist - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone :smile: --- .../workflows/python-integration-tests.yml | 26 ++ python/mypy.ini | 4 - python/poetry.lock | 20 +- python/pyproject.toml | 8 +- .../ollama_prompt_execution_settings.py | 19 +- .../connectors/ai/ollama/ollama_settings.py | 29 ++ .../ai/ollama/services/ollama_base.py | 17 + .../ollama/services/ollama_chat_completion.py | 263 ++++++++------ .../ollama/services/ollama_text_completion.py | 143 +++++--- .../ollama/services/ollama_text_embedding.py | 105 ++++-- .../connectors/ai/ollama/utils.py | 17 - .../services/open_ai_text_embedding_base.py | 2 +- .../completions/test_chat_completions.py | 21 ++ python/tests/samples/test_concepts.py | 9 +- .../tests/unit/connectors/ollama/conftest.py | 82 +++++ .../services/test_ollama_chat_completion.py | 324 +++++++++++++++--- .../services/test_ollama_text_completion.py | 169 +++++++-- .../services/test_ollama_text_embedding.py | 119 ++++++- 18 files changed, 1072 insertions(+), 305 deletions(-) create mode 100644 python/semantic_kernel/connectors/ai/ollama/ollama_settings.py create mode 100644 python/semantic_kernel/connectors/ai/ollama/services/ollama_base.py delete mode 100644 python/semantic_kernel/connectors/ai/ollama/utils.py create mode 100644 python/tests/unit/connectors/ollama/conftest.py diff --git a/.github/workflows/python-integration-tests.yml b/.github/workflows/python-integration-tests.yml index 076c66b336..125876de42 100644 --- a/.github/workflows/python-integration-tests.yml +++ b/.github/workflows/python-integration-tests.yml @@ -70,6 +70,17 @@ jobs: python -m pip install --upgrade pip setuptools wheel cd python poetry install --with tests + - name: Install Ollama + if: matrix.os == 'ubuntu-latest' + run: | + curl -fsSL https://ollama.com/install.sh | sh + ollama serve & + sleep 5 + - name: Pull model in Ollama + if: matrix.os == 'ubuntu-latest' + run: | + ollama pull ${{ vars.OLLAMA_MODEL }} + ollama list - name: Run Integration Tests id: run_tests shell: bash @@ -98,6 +109,7 @@ jobs: ACA_POOL_MANAGEMENT_ENDPOINT: ${{secrets.ACA_POOL_MANAGEMENT_ENDPOINT}} MISTRALAI_API_KEY: ${{secrets.MISTRALAI_API_KEY}} MISTRALAI_CHAT_MODEL_ID: ${{ vars.MISTRALAI_CHAT_MODEL_ID }} + OLLAMA_MODEL: "${{ matrix.os == 'ubuntu-latest' && vars.OLLAMA_MODEL || '' }}" # phi3 run: | if ${{ matrix.os == 'ubuntu-latest' }}; then docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest @@ -139,6 +151,19 @@ jobs: python -m pip install --upgrade pip setuptools wheel cd python && poetry install --with tests + - name: Install Ollama + if: matrix.os == 'ubuntu-latest' + run: | + curl -fsSL https://ollama.com/install.sh | sh + ollama serve & + sleep 5 + + - name: Pull model in Ollama + if: matrix.os == 'ubuntu-latest' + run: | + ollama pull ${{ vars.OLLAMA_MODEL }} + ollama list + - name: Run Integration Tests id: run_tests shell: bash @@ -167,6 +192,7 @@ jobs: ACA_POOL_MANAGEMENT_ENDPOINT: ${{secrets.ACA_POOL_MANAGEMENT_ENDPOINT}} MISTRALAI_API_KEY: ${{secrets.MISTRALAI_API_KEY}} MISTRALAI_CHAT_MODEL_ID: ${{ vars.MISTRALAI_CHAT_MODEL_ID }} + OLLAMA_MODEL: "${{ matrix.os == 'ubuntu-latest' && vars.OLLAMA_MODEL || '' }}" # phi3 run: | if ${{ matrix.os == 'ubuntu-latest' }}; then docker run -d --name redis-stack-server -p 6379:6379 redis/redis-stack-server:latest diff --git a/python/mypy.ini b/python/mypy.ini index 611f1b1031..6a30f83bd1 100644 --- a/python/mypy.ini +++ b/python/mypy.ini @@ -13,10 +13,6 @@ warn_untyped_fields = true [mypy-semantic_kernel] no_implicit_reexport = true -[mypy-semantic_kernel.connectors.ai.ollama.*] -ignore_errors = true -# TODO (eavanvalkenburg): remove this: https://github.com/microsoft/semantic-kernel/issues/7134 - [mypy-semantic_kernel.memory.*] ignore_errors = true # TODO (eavanvalkenburg): remove this diff --git a/python/poetry.lock b/python/poetry.lock index e7d3c431f8..8898116052 100644 --- a/python/poetry.lock +++ b/python/poetry.lock @@ -2372,7 +2372,6 @@ python-versions = ">=3.7" files = [ {file = "milvus_lite-2.4.7-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:c828190118b104b05b8c8e0b5a4147811c86b54b8fb67bc2e726ad10fc0b544e"}, {file = "milvus_lite-2.4.7-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e1537633c39879714fb15082be56a4b97f74c905a6e98e302ec01320561081af"}, - {file = "milvus_lite-2.4.7-py3-none-manylinux2014_aarch64.whl", hash = "sha256:fcb909d38c83f21478ca9cb500c84264f988c69f62715ae9462e966767fb76dd"}, {file = "milvus_lite-2.4.7-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f016474d663045787dddf1c3aad13b7d8b61fd329220318f858184918143dcbf"}, ] @@ -3142,6 +3141,20 @@ rsa = ["cryptography (>=3.0.0)"] signals = ["blinker (>=1.4.0)"] signedtoken = ["cryptography (>=3.0.0)", "pyjwt (>=2.0.0,<3)"] +[[package]] +name = "ollama" +version = "0.2.1" +description = "The official Python client for Ollama." +optional = false +python-versions = "<4.0,>=3.8" +files = [ + {file = "ollama-0.2.1-py3-none-any.whl", hash = "sha256:b6e2414921c94f573a903d1069d682ba2fb2607070ea9e19ca4a7872f2a460ec"}, + {file = "ollama-0.2.1.tar.gz", hash = "sha256:fa316baa9a81eac3beb4affb0a17deb3008fdd6ed05b123c26306cfbe4c349b6"}, +] + +[package.dependencies] +httpx = ">=0.27.0,<0.28.0" + [[package]] name = "onnxruntime" version = "1.18.0" @@ -6844,7 +6857,7 @@ doc = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linke test = ["big-O", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more-itertools", "pytest (>=6,!=8.1.*)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy", "pytest-ruff (>=0.2.1)"] [extras] -all = ["azure-ai-inference", "azure-core", "azure-cosmos", "azure-identity", "azure-search-documents", "chromadb", "ipykernel", "milvus", "mistralai", "motor", "pinecone-client", "psycopg", "pyarrow", "pymilvus", "qdrant-client", "redis", "sentence-transformers", "transformers", "usearch", "weaviate-client"] +all = ["azure-ai-inference", "azure-core", "azure-cosmos", "azure-identity", "azure-search-documents", "chromadb", "ipykernel", "milvus", "mistralai", "motor", "ollama", "pinecone-client", "psycopg", "pyarrow", "pymilvus", "qdrant-client", "redis", "sentence-transformers", "transformers", "usearch", "weaviate-client"] azure = ["azure-ai-inference", "azure-core", "azure-cosmos", "azure-identity", "azure-search-documents"] chromadb = ["chromadb"] hugging-face = ["sentence-transformers", "transformers"] @@ -6852,6 +6865,7 @@ milvus = ["milvus", "pymilvus"] mistralai = ["mistralai"] mongo = ["motor"] notebooks = ["ipykernel"] +ollama = ["ollama"] pinecone = ["pinecone-client"] postgres = ["psycopg"] qdrant = ["qdrant-client"] @@ -6862,4 +6876,4 @@ weaviate = ["weaviate-client"] [metadata] lock-version = "2.0" python-versions = "^3.10,<3.13" -content-hash = "3d6338982c9871c48bb1ed02967504967163767b0afaf50e96a1b14aa2fe0344" +content-hash = "c16c2c216ef84d091aa8f5327d75ba8a45c713f351953f53c2045bad734528c9" diff --git a/python/pyproject.toml b/python/pyproject.toml index f72ae417f3..55a8d5971f 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -54,6 +54,8 @@ pymilvus = { version = ">=2.3,<2.4.4", optional = true} milvus = { version = ">=2.3,<2.3.8", markers = 'sys_platform != "win32"', optional = true} # mistralai mistralai = { version = "^0.4.1", optional = true} +# ollama +ollama = { version = "^0.2.1", optional = true} # pinecone pinecone-client = { version = ">=3.0.0", optional = true} # postgres @@ -89,6 +91,7 @@ azure-search-documents = {version = "11.6.0b4", allow-prereleases = true} azure-core = "^1.28.0" azure-cosmos = "^4.7.0" mistralai = "^0.4.1" +ollama = "^0.2.1" transformers = { version = "^4.28.1", extras=["torch"]} sentence-transformers = "^2.2.2" @@ -113,6 +116,8 @@ pymilvus = ">=2.3,<2.4.4" milvus = { version = ">=2.3,<2.3.8", markers = 'sys_platform != "win32"'} # mistralai mistralai = "^0.4.1" +# ollama +ollama = "^0.2.1" # mongodb motor = "^3.3.2" # pinecone @@ -131,13 +136,14 @@ weaviate-client = ">=3.18,<5.0" # Extras are exposed to pip, this allows a user to easily add the right dependencies to their environment [tool.poetry.extras] -all = ["transformers", "sentence-transformers", "qdrant-client", "chromadb", "pymilvus", "milvus","mistralai", "weaviate-client", "pinecone-client", "psycopg", "redis", "azure-ai-inference", "azure-search-documents", "azure-core", "azure-identity", "azure-cosmos", "usearch", "pyarrow", "ipykernel", "motor"] +all = ["transformers", "sentence-transformers", "qdrant-client", "chromadb", "pymilvus", "milvus", "mistralai", "ollama", "weaviate-client", "pinecone-client", "psycopg", "redis", "azure-ai-inference", "azure-search-documents", "azure-core", "azure-identity", "azure-cosmos", "usearch", "pyarrow", "ipykernel", "motor"] azure = ["azure-ai-inference", "azure-search-documents", "azure-core", "azure-identity", "azure-cosmos", "msgraph-sdk"] chromadb = ["chromadb"] hugging_face = ["transformers", "sentence-transformers"] milvus = ["pymilvus", "milvus"] mistralai = ["mistralai"] +ollama = ["ollama"] mongo = ["motor"] notebooks = ["ipykernel"] pinecone = ["pinecone-client"] 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 01e8962dc1..7e365bea3d 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 @@ -2,25 +2,28 @@ from typing import Any, Literal -from pydantic import Field - from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings class OllamaPromptExecutionSettings(PromptExecutionSettings): - ai_model_id: str = Field("", serialization_alias="model") + """Settings for Ollama prompt execution.""" + format: Literal["json"] | None = None options: dict[str, Any] | None = None - stream: bool = False class OllamaTextPromptExecutionSettings(OllamaPromptExecutionSettings): - prompt: str | None = None - context: str | None = None + """Settings for Ollama text prompt execution.""" + system: str | None = None template: str | None = None - raw: bool = False + context: str | None = None + raw: bool | None = None class OllamaChatPromptExecutionSettings(OllamaPromptExecutionSettings): - messages: list[dict[str, str]] | None = None + """Settings for Ollama chat prompt execution.""" + + +class OllamaEmbeddingPromptExecutionSettings(OllamaPromptExecutionSettings): + """Settings for Ollama embedding prompt execution.""" diff --git a/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py b/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py new file mode 100644 index 0000000000..f5b3f31142 --- /dev/null +++ b/python/semantic_kernel/connectors/ai/ollama/ollama_settings.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft. All rights reserved. + +from typing import ClassVar + +from semantic_kernel.kernel_pydantic import KernelBaseSettings + + +class OllamaSettings(KernelBaseSettings): + """Ollama settings. + + The settings are first loaded from environment variables with + the prefix 'OLLAMA_'. + If the environment variables are not found, the settings can + be loaded from a .env file with the encoding 'utf-8'. + If the settings are not found in the .env file, the settings + are ignored; however, validation will fail alerting that the + settings are missing. + + Required settings for prefix 'OLLAMA' are: + - model: str - Model name. (Env var OLLAMA_MODEL) + + Optional settings for prefix 'OLLAMA' are: + - host: HttpsUrl - The endpoint of the Ollama service. (Env var OLLAMA_HOST) + """ + + env_prefix: ClassVar[str] = "OLLAMA_" + + model: str + host: str | None = None diff --git a/python/semantic_kernel/connectors/ai/ollama/services/ollama_base.py b/python/semantic_kernel/connectors/ai/ollama/services/ollama_base.py new file mode 100644 index 0000000000..ceffb48d9d --- /dev/null +++ b/python/semantic_kernel/connectors/ai/ollama/services/ollama_base.py @@ -0,0 +1,17 @@ +# Copyright (c) Microsoft. All rights reserved. + +from abc import ABC + +from ollama import AsyncClient + +from semantic_kernel.kernel_pydantic import KernelBaseModel + + +class OllamaBase(KernelBaseModel, ABC): + """Ollama service base. + + Args: + client [AsyncClient]: An Ollama client to use for the service. + """ + + client: AsyncClient diff --git a/python/semantic_kernel/connectors/ai/ollama/services/ollama_chat_completion.py b/python/semantic_kernel/connectors/ai/ollama/services/ollama_chat_completion.py index db35a54c58..1c3ffe3080 100644 --- a/python/semantic_kernel/connectors/ai/ollama/services/ollama_chat_completion.py +++ b/python/semantic_kernel/connectors/ai/ollama/services/ollama_chat_completion.py @@ -1,16 +1,22 @@ # Copyright (c) Microsoft. All rights reserved. -import json import logging -from collections.abc import AsyncGenerator -from typing import Any +import sys +from collections.abc import AsyncGenerator, AsyncIterator, Mapping +from typing import TYPE_CHECKING, Any -import aiohttp -from pydantic import HttpUrl +if sys.version_info >= (3, 12): + from typing import override # pragma: no cover +else: + from typing_extensions import override # pragma: no cover + +from ollama import AsyncClient, Message +from pydantic import ValidationError from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaChatPromptExecutionSettings -from semantic_kernel.connectors.ai.ollama.utils import AsyncSession +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_completion_client_base import TextCompletionClientBase from semantic_kernel.contents import AuthorRole from semantic_kernel.contents.chat_history import ChatHistory @@ -18,28 +24,60 @@ from semantic_kernel.contents.chat_message_content import ChatMessageContent from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent from semantic_kernel.contents.streaming_text_content import StreamingTextContent from semantic_kernel.contents.text_content import TextContent +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError + +if TYPE_CHECKING: + from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings logger: logging.Logger = logging.getLogger(__name__) -class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase): +class OllamaChatCompletion(OllamaBase, TextCompletionClientBase, ChatCompletionClientBase): """Initializes a new instance of the OllamaChatCompletion class. Make sure to have the ollama service running either locally or remotely. - - Args: - ai_model_id (str): Ollama model name, see https://ollama.ai/library - url (Optional[Union[str, HttpUrl]]): URL of the Ollama server, defaults to http://localhost:11434/api/chat - session (Optional[aiohttp.ClientSession]): Optional client session to use for requests. """ - url: HttpUrl = "http://localhost:11434/api/chat" - session: aiohttp.ClientSession | None = None + 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 OllamaChatCompletion 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.create( + model=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 + + super().__init__( + service_id=service_id or ollama_settings.model, + ai_model_id=ollama_settings.model, + client=client or AsyncClient(host=ollama_settings.host), + ) async def get_chat_message_contents( self, chat_history: ChatHistory, - settings: OllamaChatPromptExecutionSettings, + settings: "PromptExecutionSettings", **kwargs: Any, ) -> list[ChatMessageContent]: """This is the method that is called from the kernel to get a response from a chat-optimized LLM. @@ -53,29 +91,35 @@ class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase): Returns: List[ChatMessageContent]: A list of ChatMessageContent objects representing the response(s) from the LLM. """ - if not settings.ai_model_id: - settings.ai_model_id = self.ai_model_id - settings.messages = self._prepare_chat_history_for_request(chat_history) - settings.stream = False - async with ( - AsyncSession(self.session) as session, - session.post(str(self.url), json=settings.prepare_settings_dict()) as response, - ): - response.raise_for_status() - response_object = await response.json() - return [ - ChatMessageContent( - inner_content=response_object, - ai_model_id=self.ai_model_id, - role=AuthorRole.ASSISTANT, - content=response_object.get("message", {"content": None}).get("content", None), - ) - ] + settings = self.get_prompt_execution_settings_from_settings(settings) + prepared_chat_history = self._prepare_chat_history_for_request(chat_history) + + response_object = await self.client.chat( + model=self.ai_model_id, + messages=prepared_chat_history, + stream=False, + **settings.prepare_settings_dict(), + ) + + if not isinstance(response_object, Mapping): + raise ServiceInvalidResponseError( + "Invalid response type from Ollama chat completion. " + f"Expected Mapping but got {type(response_object)}." + ) + + return [ + ChatMessageContent( + inner_content=response_object, + ai_model_id=self.ai_model_id, + role=AuthorRole.ASSISTANT, + content=response_object.get("message", {"content": None}).get("content", None), + ) + ] async def get_streaming_chat_message_contents( self, chat_history: ChatHistory, - settings: OllamaChatPromptExecutionSettings, + settings: "PromptExecutionSettings", **kwargs: Any, ) -> AsyncGenerator[list[StreamingChatMessageContent], Any]: """Streams a text completion using an Ollama model. @@ -85,73 +129,81 @@ class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase): Args: chat_history (ChatHistory): A chat history that contains a list of chat messages, that can be rendered into a set of messages, from system, user, assistant and function. - settings (OllamaChatPromptExecutionSettings): Request settings. + settings (PromptExecutionSettings): Request settings. kwargs (Dict[str, Any]): The optional arguments. Yields: List[StreamingChatMessageContent]: Stream of StreamingChatMessageContent objects. """ - if not settings.ai_model_id: - settings.ai_model_id = self.ai_model_id - settings.messages = self._prepare_chat_history_for_request(chat_history) - settings.stream = True - async with ( - AsyncSession(self.session) as session, - session.post(str(self.url), json=settings.prepare_settings_dict()) as response, - ): - response.raise_for_status() - async for line in response.content: - body = json.loads(line) - if body.get("done") and body.get("message", {}).get("content") is None: - break - yield [ - StreamingChatMessageContent( - role=AuthorRole.ASSISTANT, - choice_index=0, - inner_content=body, - ai_model_id=self.ai_model_id, - content=body.get("message", {"content": None}).get("content", None), - ) - ] - if body.get("done"): - break + settings = self.get_prompt_execution_settings_from_settings(settings) + prepared_chat_history = self._prepare_chat_history_for_request(chat_history) + + response_object = await self.client.chat( + model=self.ai_model_id, + messages=prepared_chat_history, + stream=True, + **settings.prepare_settings_dict(), + ) + + if not isinstance(response_object, AsyncIterator): + raise ServiceInvalidResponseError( + "Invalid response type from Ollama chat completion. " + f"Expected AsyncIterator but got {type(response_object)}." + ) + + async for part in response_object: + yield [ + StreamingChatMessageContent( + role=AuthorRole.ASSISTANT, + choice_index=0, + inner_content=part, + ai_model_id=self.ai_model_id, + content=part.get("message", {"content": None}).get("content", None), + ) + ] async def get_text_contents( self, prompt: str, - settings: OllamaChatPromptExecutionSettings, + settings: "PromptExecutionSettings", ) -> list[TextContent]: """This is the method that is called from the kernel to get a response from a text-optimized LLM. Args: prompt (str): A prompt to complete - settings (OllamaChatPromptExecutionSettings): Settings for the request. + settings (PromptExecutionSettings): Settings for the request. Returns: List["TextContent"]: The completion result(s). """ - if not settings.ai_model_id: - settings.ai_model_id = self.ai_model_id - settings.messages = [{"role": AuthorRole.USER, "content": prompt}] - settings.stream = False - async with ( - AsyncSession(self.session) as session, - session.post(str(self.url), json=settings.prepare_settings_dict()) as response, - ): - response.raise_for_status() - response_object = await response.json() - return [ - TextContent( - inner_content=response_object, - ai_model_id=self.ai_model_id, - text=response_object.get("message", {"content": None}).get("content", None), - ) - ] + settings = self.get_prompt_execution_settings_from_settings(settings) + prepared_chat_history = [Message(role=AuthorRole.USER.value, content=prompt)] + + response_object = await self.client.chat( + model=self.ai_model_id, + messages=prepared_chat_history, + stream=False, + **settings.prepare_settings_dict(), + ) + + if not isinstance(response_object, Mapping): + raise ServiceInvalidResponseError( + "Invalid response type from Ollama chat completion. " + f"Expected Mapping but got {type(response_object)}." + ) + + return [ + TextContent( + inner_content=response_object, + ai_model_id=self.ai_model_id, + text=response_object.get("message", {"content": None}).get("content", None), + ) + ] async def get_streaming_text_contents( self, prompt: str, - settings: OllamaChatPromptExecutionSettings, + settings: "PromptExecutionSettings", ) -> AsyncGenerator[list[StreamingTextContent], Any]: """Streams a text completion using an Ollama model. @@ -159,35 +211,38 @@ class OllamaChatCompletion(TextCompletionClientBase, ChatCompletionClientBase): Args: prompt (str): A chat history that contains the prompt to complete. - settings (OllamaChatPromptExecutionSettings): Request settings. + settings (PromptExecutionSettings): Request settings. Yields: List["StreamingTextContent"]: The result stream made up of StreamingTextContent objects. """ - if not settings.ai_model_id: - settings.ai_model_id = self.ai_model_id - settings.messages = [{"role": AuthorRole.USER, "content": prompt}] - settings.stream = True - async with ( - AsyncSession(self.session) as session, - session.post(str(self.url), json=settings.prepare_settings_dict()) as response, - ): - response.raise_for_status() - async for line in response.content: - body = json.loads(line) - if body.get("done") and body.get("message", {}).get("content") is None: - break - yield [ - StreamingTextContent( - choice_index=0, - inner_content=body, - ai_model_id=self.ai_model_id, - text=body.get("message", {"content": None}).get("content", None), - ) - ] - if body.get("done"): - break + settings = self.get_prompt_execution_settings_from_settings(settings) + prepared_chat_history = [Message(role=AuthorRole.USER.value, content=prompt)] - def get_prompt_execution_settings_class(self) -> "OllamaChatPromptExecutionSettings": + response_object = await self.client.chat( + model=self.ai_model_id, + messages=prepared_chat_history, + stream=True, + **settings.prepare_settings_dict(), + ) + + if not isinstance(response_object, AsyncIterator): + raise ServiceInvalidResponseError( + "Invalid response type from Ollama chat completion. " + f"Expected AsyncIterator but got {type(response_object)}." + ) + + async for part in response_object: + yield [ + StreamingTextContent( + choice_index=0, + inner_content=part, + ai_model_id=self.ai_model_id, + text=part.get("message", {"content": None}).get("content", None), + ) + ] + + @override + def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]: """Get the request settings class.""" return OllamaChatPromptExecutionSettings diff --git a/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_completion.py b/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_completion.py index 9c1eee449e..351c4e768f 100644 --- a/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_completion.py +++ b/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_completion.py @@ -1,39 +1,78 @@ # Copyright (c) Microsoft. All rights reserved. -import json import logging -from collections.abc import AsyncGenerator -from typing import Any +import sys +from collections.abc import AsyncGenerator, AsyncIterator, Mapping +from typing import TYPE_CHECKING, Any -import aiohttp -from pydantic import HttpUrl +if sys.version_info >= (3, 12): + from typing import override # pragma: no cover +else: + from typing_extensions import override # pragma: no cover + +from ollama import AsyncClient +from pydantic import ValidationError from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaTextPromptExecutionSettings -from semantic_kernel.connectors.ai.ollama.utils import AsyncSession +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_completion_client_base import TextCompletionClientBase from semantic_kernel.contents.streaming_text_content import StreamingTextContent from semantic_kernel.contents.text_content import TextContent +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError + +if TYPE_CHECKING: + from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings logger: logging.Logger = logging.getLogger(__name__) -class OllamaTextCompletion(TextCompletionClientBase): +class OllamaTextCompletion(OllamaBase, TextCompletionClientBase): """Initializes a new instance of the OllamaTextCompletion class. Make sure to have the ollama service running either locally or remotely. - - Args: - ai_model_id (str): Ollama model name, see https://ollama.ai/library - url (Optional[Union[str, HttpUrl]]): URL of the Ollama server, defaults to http://localhost:11434/api/generate """ - url: HttpUrl = "http://localhost:11434/api/generate" - session: aiohttp.ClientSession | None = None + 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 OllamaChatCompletion 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.create( + model=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 + + super().__init__( + service_id=service_id or ollama_settings.model, + ai_model_id=ollama_settings.model, + client=client or AsyncClient(host=ollama_settings.host), + ) async def get_text_contents( self, prompt: str, - settings: OllamaTextPromptExecutionSettings, + settings: "PromptExecutionSettings", ) -> list[TextContent]: """This is the method that is called from the kernel to get a response from a text-optimized LLM. @@ -44,23 +83,29 @@ class OllamaTextCompletion(TextCompletionClientBase): Returns: List[TextContent]: A list of TextContent objects representing the response(s) from the LLM. """ - if not settings.ai_model_id: - settings.ai_model_id = self.ai_model_id - settings.prompt = prompt - settings.stream = False - async with ( - AsyncSession(self.session) as session, - session.post(self.url, json=settings.prepare_settings_dict()) as response, - ): - response.raise_for_status() - inner_content = await response.json() - text = inner_content["response"] - return [TextContent(inner_content=inner_content, ai_model_id=self.ai_model_id, text=text)] + settings = self.get_prompt_execution_settings_from_settings(settings) + + response_object = await self.client.generate( + model=self.ai_model_id, + prompt=prompt, + stream=False, + **settings.prepare_settings_dict(), + ) + + if not isinstance(response_object, Mapping): + raise ServiceInvalidResponseError( + "Invalid response type from Ollama chat completion. " + f"Expected Mapping but got {type(response_object)}." + ) + + inner_content = response_object + text = inner_content["response"] + return [TextContent(inner_content=inner_content, ai_model_id=self.ai_model_id, text=text)] async def get_streaming_text_contents( self, prompt: str, - settings: OllamaTextPromptExecutionSettings, + settings: "PromptExecutionSettings", ) -> AsyncGenerator[list[StreamingTextContent], Any]: """Streams a text completion using an Ollama model. @@ -74,27 +119,29 @@ class OllamaTextCompletion(TextCompletionClientBase): Yields: List[StreamingTextContent]: Completion result. """ - if not settings.ai_model_id: - settings.ai_model_id = self.ai_model_id - settings.prompt = prompt - settings.stream = True - async with ( - AsyncSession(self.session) as session, - session.post(self.url, json=settings.prepare_settings_dict()) as response, - ): - response.raise_for_status() - async for line in response.content: - body = json.loads(line) - if body.get("done") and body.get("response") is None: - break - yield [ - StreamingTextContent( - choice_index=0, inner_content=body, ai_model_id=self.ai_model_id, text=body.get("response") - ) - ] - if body.get("done"): - break + settings = self.get_prompt_execution_settings_from_settings(settings) - def get_prompt_execution_settings_class(self) -> "OllamaTextPromptExecutionSettings": + response_object = await self.client.generate( + model=self.ai_model_id, + prompt=prompt, + stream=True, + **settings.prepare_settings_dict(), + ) + + if not isinstance(response_object, AsyncIterator): + raise ServiceInvalidResponseError( + "Invalid response type from Ollama chat completion. " + f"Expected AsyncIterator but got {type(response_object)}." + ) + + async for part in response_object: + yield [ + StreamingTextContent( + choice_index=0, inner_content=part, ai_model_id=self.ai_model_id, text=part.get("response") + ) + ] + + @override + def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]: """Get the request settings class.""" return OllamaTextPromptExecutionSettings diff --git a/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_embedding.py b/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_embedding.py index 5716a043bb..673afc0974 100644 --- a/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_embedding.py +++ b/python/semantic_kernel/connectors/ai/ollama/services/ollama_text_embedding.py @@ -2,51 +2,108 @@ import logging import sys -from typing import Any +from typing import TYPE_CHECKING, Any + +from ollama import AsyncClient +from pydantic import ValidationError + +from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaEmbeddingPromptExecutionSettings +from semantic_kernel.connectors.ai.ollama.ollama_settings import OllamaSettings +from semantic_kernel.connectors.ai.ollama.services.ollama_base import OllamaBase +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError + +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 -import aiohttp from numpy import array, ndarray -from pydantic import HttpUrl from semantic_kernel.connectors.ai.embeddings.embedding_generator_base import EmbeddingGeneratorBase -from semantic_kernel.connectors.ai.ollama.utils import AsyncSession from semantic_kernel.utils.experimental_decorator import experimental_class logger: logging.Logger = logging.getLogger(__name__) @experimental_class -class OllamaTextEmbedding(EmbeddingGeneratorBase): +class OllamaTextEmbedding(OllamaBase, EmbeddingGeneratorBase): """Ollama embeddings client. Make sure to have the ollama service running either locally or remotely. - - Args: - ai_model_id (str): Ollama model name, see https://ollama.ai/library - url (Optional[Union[str, HttpUrl]]): URL of the Ollama server, defaults to http://localhost:11434/api/embeddings - session (Optional[aiohttp.ClientSession]): Optional client session to use for requests. """ - url: HttpUrl = "http://localhost:11434/api/embeddings" - session: aiohttp.ClientSession | None = None + 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 OllamaChatCompletion 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.create( + model=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 + + super().__init__( + service_id=service_id or ollama_settings.model, + ai_model_id=ollama_settings.model, + client=client or AsyncClient(host=ollama_settings.host), + ) @override - async def generate_embeddings(self, texts: list[str], **kwargs: Any) -> ndarray: + async def generate_embeddings( + self, + texts: list[str], + settings: "PromptExecutionSettings | None" = None, + **kwargs: Any, + ) -> ndarray: + raw_embeddings = await self.generate_raw_embeddings(texts, settings, **kwargs) + return array(raw_embeddings) + + @override + async def generate_raw_embeddings( + self, + texts: list[str], + settings: "PromptExecutionSettings | None" = None, + **kwargs: Any, + ) -> Any: + if not settings: + settings = OllamaEmbeddingPromptExecutionSettings() + else: + settings = self.get_prompt_execution_settings_from_settings(settings) + result = [] for text in texts: - async with ( - AsyncSession(self.session) as session, - session.post( - self.url, - json={"model": self.ai_model_id, "prompt": text, "options": kwargs}, - ) as response, - ): - response.raise_for_status() - response = await response.json() - result.append(response["embedding"]) - return array(result) + response_object = await self.client.embeddings( + model=self.ai_model_id, + prompt=text, + **settings.prepare_settings_dict(), + ) + result.append(response_object["embedding"]) + + return result + + @override + def get_prompt_execution_settings_class(self) -> type["PromptExecutionSettings"]: + return OllamaEmbeddingPromptExecutionSettings diff --git a/python/semantic_kernel/connectors/ai/ollama/utils.py b/python/semantic_kernel/connectors/ai/ollama/utils.py deleted file mode 100644 index 4b21cf5199..0000000000 --- a/python/semantic_kernel/connectors/ai/ollama/utils.py +++ /dev/null @@ -1,17 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import aiohttp - - -class AsyncSession: - def __init__(self, session: aiohttp.ClientSession = None): - """Initialize the AsyncSession.""" - self._session = session if session else aiohttp.ClientSession() - - async def __aenter__(self): - """Enter the context manager.""" - return await self._session.__aenter__() - - async def __aexit__(self, *args, **kwargs): - """Exit the context manager.""" - await self._session.close() diff --git a/python/semantic_kernel/connectors/ai/open_ai/services/open_ai_text_embedding_base.py b/python/semantic_kernel/connectors/ai/open_ai/services/open_ai_text_embedding_base.py index 81601912ab..73a74884ec 100644 --- a/python/semantic_kernel/connectors/ai/open_ai/services/open_ai_text_embedding_base.py +++ b/python/semantic_kernel/connectors/ai/open_ai/services/open_ai_text_embedding_base.py @@ -32,7 +32,7 @@ class OpenAITextEmbeddingBase(OpenAIHandler, EmbeddingGeneratorBase): **kwargs: Any, ) -> ndarray: raw_embeddings = await self.generate_raw_embeddings(texts, settings, batch_size, **kwargs) - return array([array(emb) for emb in raw_embeddings]) + return array(raw_embeddings) @override async def generate_raw_embeddings( diff --git a/python/tests/integration/completions/test_chat_completions.py b/python/tests/integration/completions/test_chat_completions.py index 03ac8ea8e9..4c428f524f 100644 --- a/python/tests/integration/completions/test_chat_completions.py +++ b/python/tests/integration/completions/test_chat_completions.py @@ -22,6 +22,8 @@ from semantic_kernel.connectors.ai.mistral_ai.prompt_execution_settings.mistral_ MistralAIChatPromptExecutionSettings, ) from semantic_kernel.connectors.ai.mistral_ai.services.mistral_ai_chat_completion import MistralAIChatCompletion +from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaChatPromptExecutionSettings +from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import ( AzureChatPromptExecutionSettings, ) @@ -47,6 +49,13 @@ try: except KeyError: mistral_ai_setup = False +ollama_setup: bool = False +try: + if os.environ["OLLAMA_MODEL"]: + ollama_setup = True +except KeyError: + ollama_setup = False + def setup( kernel: Kernel, @@ -101,6 +110,7 @@ def services() -> dict[str, tuple[ChatCompletionClientBase, type[PromptExecution "azure_custom_client": (azure_custom_client, AzureChatPromptExecutionSettings), "azure_ai_inference": (azure_ai_inference_client, AzureAIInferenceChatPromptExecutionSettings), "mistral_ai": (MistralAIChatCompletion() if mistral_ai_setup else None, MistralAIChatPromptExecutionSettings), + "ollama": (OllamaChatCompletion(), OllamaChatPromptExecutionSettings), } @@ -432,6 +442,17 @@ pytestmark = pytest.mark.parametrize( marks=pytest.mark.skipif(not mistral_ai_setup, reason="Mistral AI Environment Variables not set"), id="mistral_ai_text_input", ), + pytest.param( + "ollama", + {}, + [ + ChatMessageContent(role=AuthorRole.USER, items=[TextContent(text="Hello")]), + ChatMessageContent(role=AuthorRole.USER, items=[TextContent(text="How are you today?")]), + ], + ["Hello", "well"], + marks=pytest.mark.skipif(not ollama_setup, reason="Need local Ollama setup"), + id="ollama_text_input", + ), ], ) diff --git a/python/tests/samples/test_concepts.py b/python/tests/samples/test_concepts.py index 32a505926e..64f6f2c682 100644 --- a/python/tests/samples/test_concepts.py +++ b/python/tests/samples/test_concepts.py @@ -84,7 +84,14 @@ concepts = [ ["Create a secret with the name 'Foo' and value 'Bar'", "exit"], id="openai_plugin_azure_key_vault", ), - param(openai_plugin_klarna, [], id="openai_plugin_klarna"), + param( + openai_plugin_klarna, + [], + id="openai_plugin_klarna", + marks=pytest.mark.skip( + reason="Temporarily: https://www.klarna.com/us/shopping/public/openai/v0/api-docs/ returns 404" + ), + ), param(plugins_from_dir, [], id="plugins_from_dir"), param(azure_chat_gpt_api_handlebars, ["What is 3+3?", "exit"], id="azure_chat_gpt_api_handlebars"), param(azure_chat_gpt_api_jinja2, ["What is 3+3?", "exit"], id="azure_chat_gpt_api_jinja2"), diff --git a/python/tests/unit/connectors/ollama/conftest.py b/python/tests/unit/connectors/ollama/conftest.py new file mode 100644 index 0000000000..3d278bc026 --- /dev/null +++ b/python/tests/unit/connectors/ollama/conftest.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft. All rights reserved. + +from collections.abc import AsyncGenerator, AsyncIterator +from unittest.mock import MagicMock + +import pytest +from ollama import AsyncClient + +from semantic_kernel.contents.chat_history import ChatHistory + + +@pytest.fixture() +def model_id() -> str: + return "test_model_id" + + +@pytest.fixture() +def service_id() -> str: + return "test_service_id" + + +@pytest.fixture() +def host() -> str: + return "http://localhost:5000" + + +@pytest.fixture() +def custom_client() -> AsyncClient: + return AsyncClient("http://localhost:5001") + + +@pytest.fixture() +def chat_history() -> ChatHistory: + chat_history = ChatHistory() + chat_history.add_user_message("test_prompt") + return chat_history + + +@pytest.fixture() +def prompt() -> str: + return "test_prompt" + + +@pytest.fixture() +def default_options() -> dict: + return {"test": "test"} + + +@pytest.fixture() +def ollama_unit_test_env(monkeypatch, model_id, host, exclude_list): + """Fixture to set environment variables for OllamaSettings.""" + if exclude_list is None: + exclude_list = [] + + env_vars = { + "OLLAMA_MODEL": model_id, + "OLLAMA_HOST": host, + } + + for key, value in env_vars.items(): + if key not in exclude_list: + monkeypatch.setenv(key, value) + else: + monkeypatch.delenv(key, raising=False) + + return env_vars + + +@pytest.fixture() +def mock_streaming_text_response() -> AsyncIterator: + streaming_text_response = MagicMock(spec=AsyncGenerator) + streaming_text_response.__aiter__.return_value = [{"response": "test_response"}] + + return streaming_text_response + + +@pytest.fixture() +def mock_streaming_chat_response() -> AsyncIterator: + streaming_chat_response = MagicMock(spec=AsyncGenerator) + streaming_chat_response.__aiter__.return_value = [{"message": {"content": "test_response"}}] + + return streaming_chat_response diff --git a/python/tests/unit/connectors/ollama/services/test_ollama_chat_completion.py b/python/tests/unit/connectors/ollama/services/test_ollama_chat_completion.py index c83931ecf6..7e78e6d8a3 100644 --- a/python/tests/unit/connectors/ollama/services/test_ollama_chat_completion.py +++ b/python/tests/unit/connectors/ollama/services/test_ollama_chat_completion.py @@ -4,96 +4,310 @@ from unittest.mock import patch import pytest -from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaChatPromptExecutionSettings +from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import ( + OllamaChatPromptExecutionSettings, + OllamaTextPromptExecutionSettings, +) from semantic_kernel.connectors.ai.ollama.services.ollama_chat_completion import OllamaChatCompletion -from semantic_kernel.contents.chat_history import ChatHistory -from tests.unit.connectors.ollama.utils import MockResponse +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError -def test_settings(): - ollama = OllamaChatCompletion(ai_model_id="test_model") +def test_settings(model_id): + """Test that the settings class is correct.""" + ollama = OllamaChatCompletion(ai_model_id=model_id) settings = ollama.get_prompt_execution_settings_class() assert settings == OllamaChatPromptExecutionSettings +def test_init_empty_service_id(model_id): + """Test that the service initializes correctly with an empty service id.""" + ollama = OllamaChatCompletion(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 = OllamaChatCompletion(ai_model_id=model_id, client=custom_client) + assert ollama.client == custom_client + + +@pytest.mark.parametrize("exclude_list", [["OLLAMA_MODEL"]], 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): + _ = OllamaChatCompletion(env_file_path="fake_env_file_path.env") + + @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_complete_chat(mock_post): - mock_post.return_value = MockResponse(response={"message": {"content": "test_response"}}) - ollama = OllamaChatCompletion(ai_model_id="test_model") - chat_history = ChatHistory() - chat_history.add_user_message("test_prompt") +@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client +@patch("ollama.AsyncClient.chat") # mock_chat_client +async def test_custom_host( + mock_chat_client, + mock_client, + model_id, + service_id, + host, + chat_history, + prompt, + default_options, +): + """Test that the service initializes and generates content correctly with a custom host.""" + mock_chat_client.return_value = {"message": {"content": "test_response"}} + + ollama = OllamaChatCompletion(ai_model_id=model_id, host=host) + + chat_responses = await ollama.get_chat_message_contents( + chat_history, + OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options), + ) + + text_responses = await ollama.get_text_contents( + prompt, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ) + + # Check that the client was initialized once with the correct host + assert mock_client.call_count == 1 + mock_client.assert_called_with(host=host) + # Check that the chat client was called twice and the responses are correct + assert mock_chat_client.call_count == 2 + assert len(chat_responses) == 1 + assert chat_responses[0].content == "test_response" + assert len(text_responses) == 1 + assert text_responses[0].text == "test_response" + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client +@patch("ollama.AsyncClient.chat") # mock_chat_client +async def test_custom_host_streaming( + mock_chat_client, + mock_client, + mock_streaming_chat_response, + model_id, + service_id, + host, + chat_history, + prompt, + default_options, +): + """Test that the service initializes and generates streaming content correctly with a custom host.""" + mock_chat_client.return_value = mock_streaming_chat_response + + ollama = OllamaChatCompletion(ai_model_id=model_id, host=host) + + async for messages in ollama.get_streaming_chat_message_contents( + chat_history, + OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options), + ): + assert len(messages) == 1 + assert messages[0].role == "assistant" + assert messages[0].content == "test_response" + + async for messages in ollama.get_streaming_text_contents( + prompt, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ): + assert len(messages) == 1 + assert messages[0].text == "test_response" + + # Check that the client was initialized once with the correct host + assert mock_client.call_count == 1 + mock_client.assert_called_with(host=host) + # Check that the chat client was called twice and the responses are correct + assert mock_chat_client.call_count == 2 + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.chat") +async def test_chat_completion(mock_chat_client, model_id, service_id, chat_history, default_options): + """Test that the chat completion service completes correctly.""" + mock_chat_client.return_value = {"message": {"content": "test_response"}} + + ollama = OllamaChatCompletion(ai_model_id=model_id) response = await ollama.get_chat_message_contents( chat_history, - OllamaChatPromptExecutionSettings(service_id="test_model", ai_model_id="test_model", options={"test": "test"}), + OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options), ) + + assert len(response) == 1 assert response[0].content == "test_response" - mock_post.assert_called_once_with( - "http://localhost:11434/api/chat", - json={ - "model": "test_model", - "messages": [{"role": "user", "content": "test_prompt"}], - "options": {"test": "test"}, - "stream": False, - }, + mock_chat_client.assert_called_once_with( + model=model_id, + messages=ollama._prepare_chat_history_for_request(chat_history), + options=default_options, + stream=False, ) @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_complete(mock_post): - mock_post.return_value = MockResponse(response={"message": {"content": "test_response"}}) - ollama = OllamaChatCompletion(ai_model_id="test_model") +@patch("ollama.AsyncClient.chat") +async def test_chat_completion_wrong_return_type( + mock_chat_client, + mock_streaming_chat_response, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the chat completion service fails when the return type is incorrect.""" + mock_chat_client.return_value = mock_streaming_chat_response # should not be a streaming response + + ollama = OllamaChatCompletion(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + await ollama.get_chat_message_contents( + chat_history, + OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options), + ) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.chat") +async def test_text_completion(mock_chat_client, model_id, service_id, prompt, default_options): + """Test that the text completion service completes correctly.""" + mock_chat_client.return_value = {"message": {"content": "test_response"}} + ollama = OllamaChatCompletion(ai_model_id=model_id) response = await ollama.get_text_contents( - "test_prompt", - OllamaChatPromptExecutionSettings(service_id="test_model", ai_model_id="test_model", options={"test": "test"}), + prompt, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), ) + + assert len(response) == 1 assert response[0].text == "test_response" + mock_chat_client.assert_called_once_with( + model=model_id, + messages=[{"role": "user", "content": prompt}], + options=default_options, + stream=False, + ) @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_complete_chat_stream(mock_post): - mock_post.return_value = MockResponse(response={"message": {"content": "test_response"}}) - ollama = OllamaChatCompletion(ai_model_id="test_model") - chat_history = ChatHistory() - chat_history.add_user_message("test_prompt") +@patch("ollama.AsyncClient.chat") +async def test_text_completion_wrong_return_type( + mock_chat_client, + mock_streaming_chat_response, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the text completion service fails when the return type is incorrect.""" + mock_chat_client.return_value = mock_streaming_chat_response # should not be a streaming response + + ollama = OllamaChatCompletion(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + await ollama.get_text_contents( + chat_history, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.chat") +async def test_streaming_chat_completion( + mock_chat_client, + mock_streaming_chat_response, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the streaming chat completion service completes correctly.""" + mock_chat_client.return_value = mock_streaming_chat_response + + ollama = OllamaChatCompletion(ai_model_id=model_id) response = ollama.get_streaming_chat_message_contents( chat_history, - OllamaChatPromptExecutionSettings(ai_model_id="test_model", options={"test": "test"}), + OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options), ) + + responses = [] async for line in response: if line: assert line[0].content == "test_response" - mock_post.assert_called_once_with( - "http://localhost:11434/api/chat", - json={ - "model": "test_model", - "messages": [{"role": "user", "content": "test_prompt"}], - "options": {"test": "test"}, - "stream": True, - }, + responses.append(line[0].content) + assert len(responses) == 1 + + mock_chat_client.assert_called_once_with( + model=model_id, + messages=ollama._prepare_chat_history_for_request(chat_history), + options=default_options, + stream=True, ) @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_complete_stream(mock_post): - mock_post.return_value = MockResponse(response={"message": {"content": "test_response"}}) - ollama = OllamaChatCompletion(ai_model_id="test_model") +@patch("ollama.AsyncClient.chat") +async def test_streaming_chat_completion_wrong_return_type( + mock_chat_client, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the chat completion streaming service fails when the return type is incorrect.""" + mock_chat_client.return_value = {"message": {"content": "test_response"}} # should not be a non-streaming response + + ollama = OllamaChatCompletion(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + async for _ in ollama.get_streaming_chat_message_contents( + chat_history, + OllamaChatPromptExecutionSettings(service_id=service_id, options=default_options), + ): + pass + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.chat") +async def test_streaming_text_completion( + mock_chat_client, + mock_streaming_chat_response, + model_id, + service_id, + prompt, + default_options, +): + """Test that the streaming text completion service completes correctly.""" + mock_chat_client.return_value = mock_streaming_chat_response + + ollama = OllamaChatCompletion(ai_model_id=model_id) response = ollama.get_streaming_text_contents( - "test_prompt", - OllamaChatPromptExecutionSettings(ai_model_id="test_model", options={"test": "test"}), + prompt, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), ) + + responses = [] async for line in response: if line: assert line[0].text == "test_response" - mock_post.assert_called_once_with( - "http://localhost:11434/api/chat", - json={ - "model": "test_model", - "options": {"test": "test"}, - "stream": True, - "messages": [{"role": "user", "content": "test_prompt"}], - }, + responses.append(line[0].text) + assert len(responses) == 1 + + mock_chat_client.assert_called_once_with( + model=model_id, + messages=[{"role": "user", "content": prompt}], + options=default_options, + stream=True, ) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.chat") +async def test_streaming_text_completion_wrong_return_type( + mock_chat_client, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the text completion streaming service fails when the return type is incorrect.""" + mock_chat_client.return_value = {"message": {"content": "test_response"}} # should not be a non-streaming response + + ollama = OllamaChatCompletion(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + async for _ in ollama.get_streaming_text_contents( + chat_history, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ): + pass diff --git a/python/tests/unit/connectors/ollama/services/test_ollama_text_completion.py b/python/tests/unit/connectors/ollama/services/test_ollama_text_completion.py index 772a3a13e7..3992517d5c 100644 --- a/python/tests/unit/connectors/ollama/services/test_ollama_text_completion.py +++ b/python/tests/unit/connectors/ollama/services/test_ollama_text_completion.py @@ -6,46 +6,163 @@ import pytest from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaTextPromptExecutionSettings from semantic_kernel.connectors.ai.ollama.services.ollama_text_completion import OllamaTextCompletion -from tests.unit.connectors.ollama.utils import MockResponse +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError -def test_settings(): - ollama = OllamaTextCompletion(ai_model_id="test_model") +def test_settings(model_id): + """Test that the settings class is correct""" + ollama = OllamaTextCompletion(ai_model_id=model_id) settings = ollama.get_prompt_execution_settings_class() assert settings == OllamaTextPromptExecutionSettings +def test_init_empty_service_id(model_id): + """Test that the service initializes correctly with an empty service_id""" + ollama = OllamaTextCompletion(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 = OllamaTextCompletion(ai_model_id=model_id, client=custom_client) + assert ollama.client == custom_client + + +@pytest.mark.parametrize("exclude_list", [["OLLAMA_MODEL"]], 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): + _ = OllamaTextCompletion(env_file_path="fake_env_file_path.env") + + @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_complete(mock_post): - mock_post.return_value = MockResponse(response={"response": "test_response"}) - ollama = OllamaTextCompletion(ai_model_id="test_model") +@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client +@patch("ollama.AsyncClient.generate") # mock_completion_client +async def test_custom_host( + mock_completion_client, mock_client, model_id, service_id, host, chat_history, default_options +): + """Test that the service initializes and generates content correctly with a custom host.""" + mock_completion_client.return_value = {"response": "test_response"} + + ollama = OllamaTextCompletion(ai_model_id=model_id, host=host) + _ = await ollama.get_text_contents( + chat_history, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ) + + mock_client.assert_called_once_with(host=host) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client +@patch("ollama.AsyncClient.generate") # mock_completion_client +async def test_custom_host_streaming( + mock_completion_client, mock_client, model_id, service_id, host, chat_history, default_options +): + """Test that the service initializes and generates streaming content correctly with a custom host.""" + mock_completion_client.__aiter__.return_value = {"response": "test_response"} + + ollama = OllamaTextCompletion(ai_model_id=model_id, host=host) + async for _ in ollama.get_streaming_text_contents( + chat_history, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ): + pass + + mock_client.assert_called_once_with(host=host) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.generate") +async def test_completion(mock_completion_client, model_id, service_id, prompt, default_options): + """Test that the service generates content correctly.""" + mock_completion_client.return_value = {"response": "test_response"} + + ollama = OllamaTextCompletion(ai_model_id=model_id) response = await ollama.get_text_contents( - "test prompt", - OllamaTextPromptExecutionSettings(options={"test": "test"}), + prompt, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), ) + assert response[0].text == "test_response" + mock_completion_client.assert_called_once_with( + model=model_id, + prompt=prompt, + options=default_options, + stream=False, + ) @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_complete_stream(mock_post): - mock_post.return_value = MockResponse(response={"response": "test_response"}) - ollama = OllamaTextCompletion(ai_model_id="test_model") +@patch("ollama.AsyncClient.generate") +async def test_completion_wrong_return_type( + mock_completion_client, + mock_streaming_text_response, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the completion service fails when the return type is incorrect.""" + mock_completion_client.return_value = mock_streaming_text_response # should not be a streaming response + + ollama = OllamaTextCompletion(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + await ollama.get_text_contents( + chat_history, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.generate") +async def test_streaming_completion( + mock_completion_client, + mock_streaming_text_response, + model_id, + service_id, + prompt, + default_options, +): + """Test that the service generates streaming content correctly.""" + mock_completion_client.return_value = mock_streaming_text_response + + ollama = OllamaTextCompletion(ai_model_id=model_id) response = ollama.get_streaming_text_contents( - "test_prompt", - OllamaTextPromptExecutionSettings(options={"test": "test"}), + prompt, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), ) + + responses = [] async for line in response: - if line: - assert line[0].text == "test_response" - mock_post.assert_called_once_with( - "http://localhost:11434/api/generate", - json={ - "model": "test_model", - "options": {"test": "test"}, - "stream": True, - "prompt": "test_prompt", - "raw": False, - }, + assert line[0].text == "test_response" + responses.append(line) + assert len(responses) == 1 + + mock_completion_client.assert_called_once_with( + model=model_id, + prompt=prompt, + options=default_options, + stream=True, ) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.generate") +async def test_streaming_completion_wrong_return_type( + mock_completion_client, + model_id, + service_id, + chat_history, + default_options, +): + """Test that the streaming completion service fails when the return type is incorrect.""" + mock_completion_client.return_value = {"response": "test_response"} # should not be a non-streaming response + + ollama = OllamaTextCompletion(ai_model_id=model_id) + with pytest.raises(ServiceInvalidResponseError): + async for _ in ollama.get_streaming_text_contents( + chat_history, + OllamaTextPromptExecutionSettings(service_id=service_id, options=default_options), + ): + pass diff --git a/python/tests/unit/connectors/ollama/services/test_ollama_text_embedding.py b/python/tests/unit/connectors/ollama/services/test_ollama_text_embedding.py index 536fc70839..d008182f34 100644 --- a/python/tests/unit/connectors/ollama/services/test_ollama_text_embedding.py +++ b/python/tests/unit/connectors/ollama/services/test_ollama_text_embedding.py @@ -2,27 +2,120 @@ from unittest.mock import patch +import numpy import pytest from numpy import array +from semantic_kernel.connectors.ai.ollama.ollama_prompt_execution_settings import OllamaEmbeddingPromptExecutionSettings from semantic_kernel.connectors.ai.ollama.services.ollama_text_embedding import OllamaTextEmbedding -from tests.unit.connectors.ollama.utils import MockResponse +from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError + + +def test_init_empty_service_id(model_id): + """Test that the service initializes correctly with an empty service id.""" + ollama = OllamaTextEmbedding(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 = OllamaTextEmbedding(ai_model_id=model_id, client=custom_client) + assert ollama.client == custom_client + + +@pytest.mark.parametrize("exclude_list", [["OLLAMA_MODEL"]], 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): + _ = OllamaTextEmbedding(env_file_path="fake_env_file_path.env") @pytest.mark.asyncio -@patch("aiohttp.ClientSession.post") -async def test_embedding(mock_post): - mock_post.return_value = MockResponse(response={"embedding": [0.1, 0.2, 0.3]}) - ollama = OllamaTextEmbedding(ai_model_id="test_model") +@patch("ollama.AsyncClient.__init__", return_value=None) # mock_client +@patch("ollama.AsyncClient.embeddings") # mock_embedding_client +async def test_custom_host(mock_embedding_client, mock_client, model_id, host, prompt): + """Test that the service initializes and generates embeddings correctly with a custom host.""" + mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]} + + ollama = OllamaTextEmbedding(ai_model_id=model_id, host=host) + _ = await ollama.generate_embeddings( + [prompt], + ) + + mock_client.assert_called_once_with(host=host) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.embeddings") +async def test_embedding(mock_embedding_client, model_id, prompt): + """Test that the service initializes and generates embeddings correctly.""" + mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]} + settings = OllamaEmbeddingPromptExecutionSettings() + settings.options = {"test_key": "test_value"} + + ollama = OllamaTextEmbedding(ai_model_id=model_id) response = await ollama.generate_embeddings( - ["test_prompt"], + [prompt], + settings=settings, ) + assert response.all() == array([0.1, 0.2, 0.3]).all() - mock_post.assert_called_once_with( - "http://localhost:11434/api/embeddings", - json={ - "model": "test_model", - "prompt": "test_prompt", - "options": {}, - }, + mock_embedding_client.assert_called_once_with(model=model_id, prompt=prompt, options=settings.options) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.embeddings") +async def test_embedding_list_input(mock_embedding_client, model_id, prompt): + """Test that the service initializes and generates embeddings correctly with a list of prompts.""" + mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]} + settings = OllamaEmbeddingPromptExecutionSettings() + settings.options = {"test_key": "test_value"} + + ollama = OllamaTextEmbedding(ai_model_id=model_id) + responses = await ollama.generate_embeddings( + [prompt, prompt], + settings=settings, ) + + assert len(responses) == 2 + assert type(responses) is numpy.ndarray + assert all(type(response) is numpy.ndarray for response in responses) + assert mock_embedding_client.call_count == 2 + mock_embedding_client.assert_called_with(model=model_id, prompt=prompt, options=settings.options) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.embeddings") +async def test_raw_embedding(mock_embedding_client, model_id, prompt): + """Test that the service initializes and generates embeddings correctly.""" + mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]} + settings = OllamaEmbeddingPromptExecutionSettings() + settings.options = {"test_key": "test_value"} + + ollama = OllamaTextEmbedding(ai_model_id=model_id) + response = await ollama.generate_raw_embeddings( + [prompt], + settings=settings, + ) + + assert response == [[0.1, 0.2, 0.3]] + mock_embedding_client.assert_called_once_with(model=model_id, prompt=prompt, options=settings.options) + + +@pytest.mark.asyncio +@patch("ollama.AsyncClient.embeddings") +async def test_raw_embedding_list_input(mock_embedding_client, model_id, prompt): + """Test that the service initializes and generates embeddings correctly with a list of prompts.""" + mock_embedding_client.return_value = {"embedding": [0.1, 0.2, 0.3]} + settings = OllamaEmbeddingPromptExecutionSettings() + settings.options = {"test_key": "test_value"} + + ollama = OllamaTextEmbedding(ai_model_id=model_id) + responses = await ollama.generate_raw_embeddings( + [prompt, prompt], + settings=settings, + ) + + assert responses == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]] + assert mock_embedding_client.call_count == 2 + mock_embedding_client.assert_called_with(model=model_id, prompt=prompt, options=settings.options)