feat: Ollama - accept str as ChatGenerator input; deprecate generator; rm generator example (#3388)

Author: anakin87

integrations/ollama/examples/generator_example.py

@@ -27,7 +27,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.22.0", "ollama>=0.5.4", "pydantic>=2.12.0", "tenacity>=8.2.3"]
+dependencies = ["haystack-ai>=2.30.0", "ollama>=0.5.4", "pydantic>=2.12.0", "tenacity>=8.2.3"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/ollama#readme"

integrations/ollama/pyproject.toml

@@ -4,6 +4,7 @@
 from typing import Any, Literal
 
 from haystack import component, default_from_dict, default_to_dict
+from haystack.components.generators.utils import _normalize_messages
 from haystack.dataclasses import (
     AsyncStreamingCallbackT,
     ChatMessage,
@@ -577,7 +578,7 @@ async def _chat_async(
     @component.output_types(replies=list[ChatMessage])
     def run(
         self,
-        messages: list[ChatMessage],
+        messages: list[ChatMessage] | str,
         generation_kwargs: dict[str, Any] | None = None,
         tools: ToolsType | None = None,
         *,
@@ -587,7 +588,8 @@ def run(
         Runs an Ollama Model on a given chat history.
 
         :param messages:
-            A list of ChatMessage instances representing the input messages.
+            A list of ChatMessage instances representing the input messages. If a string is provided, it is converted
+            to a list containing a ChatMessage with user role.
         :param generation_kwargs:
             Per-call overrides for Ollama inference options.
             These are merged on top of the instance-level `generation_kwargs`.
@@ -603,6 +605,7 @@ def run(
         :returns: A dictionary with the following keys:
             - `replies`: A list of ChatMessages containing the model's response
         """
+        messages = _normalize_messages(messages)
 
         # Validate and select the streaming callback
         callback = select_streaming_callback(
@@ -636,7 +639,7 @@ def run(
     @component.output_types(replies=list[ChatMessage])
     async def run_async(
         self,
-        messages: list[ChatMessage],
+        messages: list[ChatMessage] | str,
         generation_kwargs: dict[str, Any] | None = None,
         tools: ToolsType | None = None,
         *,
@@ -646,7 +649,8 @@ async def run_async(
         Async version of run. Runs an Ollama Model on a given chat history.
 
         :param messages:
-            A list of ChatMessage instances representing the input messages.
+            A list of ChatMessage instances representing the input messages. If a string is provided, it is converted
+            to a list containing a ChatMessage with user role.
         :param generation_kwargs:
             Per-call overrides for Ollama inference options.
             These are merged on top of the instance-level `generation_kwargs`.
@@ -659,6 +663,8 @@ async def run_async(
         :returns: A dictionary with the following keys:
             - `replies`: A list of ChatMessages containing the model's response
         """
+        messages = _normalize_messages(messages)
+
         # Validate and select the streaming callback
         callback = select_streaming_callback(self.streaming_callback, streaming_callback, requires_async=True)
 

integrations/ollama/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py

@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Callable
 from typing import Any
 
@@ -139,6 +140,13 @@ def __init__(
             - any negative number which will keep the model loaded in memory (e.g. -1 or "-1m")
             - '0' which will unload the model immediately after generating a response.
         """
+        warnings.warn(
+            "The `OllamaGenerator` component is deprecated and will be removed in a future version. "
+            "Use `OllamaChatGenerator` instead, which now also supports string inputs.",
+            FutureWarning,
+            stacklevel=2,
+        )
+
         self.timeout = timeout
         self.raw = raw
         self.template = template

integrations/ollama/src/haystack_integrations/components/generators/ollama/generator.py

@@ -1007,6 +1007,57 @@ def test_run(self, mock_client):
         assert result["replies"][0].text == "Fine. How can I help you today?"
         assert result["replies"][0].role == "assistant"
 
+    @patch("haystack_integrations.components.generators.ollama.chat.chat_generator.Client")
+    def test_run_with_string_input(self, mock_client):
+        generator = OllamaChatGenerator()
+
+        mock_response = ChatResponse(
+            model="qwen3:0.6b",
+            created_at="2023-12-12T14:13:43.416799Z",
+            message={"role": "assistant", "content": "Paris"},
+            done=True,
+            prompt_eval_count=1,
+            eval_count=1,
+        )
+
+        mock_client_instance = mock_client.return_value
+        mock_client_instance.chat.return_value = mock_response
+
+        result = generator.run("What's the capital of France?")
+
+        _, kwargs = mock_client_instance.chat.call_args
+        assert kwargs["messages"] == [{"role": "user", "content": "What's the capital of France?"}]
+
+        assert isinstance(result["replies"], list)
+        assert len(result["replies"]) == 1
+        assert isinstance(result["replies"][0], ChatMessage)
+
+    @pytest.mark.asyncio
+    @patch("haystack_integrations.components.generators.ollama.chat.chat_generator.AsyncClient")
+    async def test_run_async_with_string_input(self, mock_async_client):
+        generator = OllamaChatGenerator()
+
+        mock_response = ChatResponse(
+            model="qwen3:0.6b",
+            created_at="2023-12-12T14:13:43.416799Z",
+            message={"role": "assistant", "content": "Paris"},
+            done=True,
+            prompt_eval_count=1,
+            eval_count=1,
+        )
+
+        mock_async_client_instance = mock_async_client.return_value
+        mock_async_client_instance.chat = AsyncMock(return_value=mock_response)
+
+        result = await generator.run_async("What's the capital of France?")
+
+        _, kwargs = mock_async_client_instance.chat.call_args
+        assert kwargs["messages"] == [{"role": "user", "content": "What's the capital of France?"}]
+
+        assert isinstance(result["replies"], list)
+        assert len(result["replies"]) == 1
+        assert isinstance(result["replies"][0], ChatMessage)
+
     @patch("haystack_integrations.components.generators.ollama.chat.chat_generator.Client")
     def test_run_retries_after_failure(self, mock_client):
         generator = OllamaChatGenerator(max_retries=1)

integrations/ollama/tests/test_chat_generator.py

@@ -2,3 +2,7 @@
 # attacks via compromised dependencies. uv resolves this relative to the current clock at
 # install/lock time, so no manual date updates are needed.
 exclude-newer = "24 hours"
+
+# haystack-ai is a first-party dependency
+[exclude-newer-package]
+haystack-ai = false