ollama-haystack 2.1.2__tar.gz → 2.2.0__tar.gz

{ollama_haystack-2.1.2 → ollama_haystack-2.2.0}/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [integrations/ollama-v2.1.2] - 2024-12-18
+
+### 🐛 Bug Fixes
+
+- Make Ollama Chat Generator compatible with new ChatMessage (#1256)
+
+
 ## [integrations/ollama-v2.1.1] - 2024-12-10
 
 ### 🌀 Miscellaneous

{ollama_haystack-2.1.2 → ollama_haystack-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ollama-haystack
-Version: 2.1.2
+Version: 2.2.0
 Summary: An integration between the Ollama LLM framework and Haystack
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/ollama#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues

{ollama_haystack-2.1.2 → ollama_haystack-2.2.0}/pyproject.toml

@@ -51,6 +51,7 @@ dependencies = [
     "pytest",
     "pytest-rerunfailures",
     "haystack-pydoc-tools",
+    "jsonschema",  # needed for Tool
 ]
 [tool.hatch.envs.default.scripts]
 test = "pytest {args:tests}"

{ollama_haystack-2.1.2 → ollama_haystack-2.2.0}/src/haystack_integrations/components/generators/ollama/chat/chat_generator.py

@@ -1,12 +1,68 @@
 from typing import Any, Callable, Dict, List, Optional, Union
 
 from haystack import component, default_from_dict, default_to_dict
-from haystack.dataclasses import ChatMessage, StreamingChunk
+from haystack.dataclasses import ChatMessage, StreamingChunk, ToolCall
+from haystack.tools import Tool, _check_duplicate_tool_names, deserialize_tools_inplace
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 
 from ollama import ChatResponse, Client
 
 
+def _convert_chatmessage_to_ollama_format(message: ChatMessage) -> Dict[str, Any]:
+    """
+    Convert a ChatMessage to the format expected by Ollama Chat API.
+    """
+    text_contents = message.texts
+    tool_calls = message.tool_calls
+    tool_call_results = message.tool_call_results
+
+    if not text_contents and not tool_calls and not tool_call_results:
+        msg = "A `ChatMessage` must contain at least one `TextContent`, `ToolCall`, or `ToolCallResult`."
+        raise ValueError(msg)
+    elif len(text_contents) + len(tool_call_results) > 1:
+        msg = "A `ChatMessage` can only contain one `TextContent` or one `ToolCallResult`."
+        raise ValueError(msg)
+
+    ollama_msg: Dict[str, Any] = {"role": message._role.value}
+
+    if tool_call_results:
+        # Ollama does not provide a way to communicate errors in tool invocations, so we ignore the error field
+        ollama_msg["content"] = tool_call_results[0].result
+        return ollama_msg
+
+    if text_contents:
+        ollama_msg["content"] = text_contents[0]
+    if tool_calls:
+        # Ollama does not support tool call id, so we ignore it
+        ollama_msg["tool_calls"] = [
+            {"type": "function", "function": {"name": tc.tool_name, "arguments": tc.arguments}} for tc in tool_calls
+        ]
+    return ollama_msg
+
+
+def _convert_ollama_response_to_chatmessage(ollama_response: "ChatResponse") -> ChatMessage:
+    """
+    Converts the non-streaming response from the Ollama API to a ChatMessage with assistant role.
+    """
+    response_dict = ollama_response.model_dump()
+
+    ollama_message = response_dict["message"]
+
+    text = ollama_message["content"]
+
+    tool_calls = []
+    if ollama_tool_calls := ollama_message.get("tool_calls"):
+        for ollama_tc in ollama_tool_calls:
+            tool_calls.append(
+                ToolCall(tool_name=ollama_tc["function"]["name"], arguments=ollama_tc["function"]["arguments"])
+            )
+
+    message = ChatMessage.from_assistant(text=text, tool_calls=tool_calls)
+
+    message.meta.update({key: value for key, value in response_dict.items() if key != "message"})
+    return message
+
+
 @component
 class OllamaChatGenerator:
     """
@@ -40,6 +96,7 @@ class OllamaChatGenerator:
         timeout: int = 120,
         keep_alive: Optional[Union[float, str]] = None,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
+        tools: Optional[List[Tool]] = None,
     ):
         """
         :param model:
@@ -52,9 +109,6 @@ class OllamaChatGenerator:
             [Ollama docs](https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values).
         :param timeout:
             The number of seconds before throwing a timeout error from the Ollama API.
-        :param streaming_callback:
-            A callback function that is called when a new token is received from the stream.
-            The callback function accepts StreamingChunk as an argument.
         :param keep_alive:
             The option that controls how long the model will stay loaded into memory following the request.
             If not set, it will use the default value from the Ollama (5 minutes).
@@ -63,14 +117,24 @@ class OllamaChatGenerator:
             - a number in seconds (such as 3600)
             - 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.
+        :param streaming_callback:
+            A callback function that is called when a new token is received from the stream.
+            The callback function accepts StreamingChunk as an argument.
+        :param tools:
+            A list of tools for which the model can prepare calls.
+            Not all models support tools. For a list of models compatible with tools, see the
+            [models page](https://ollama.com/search?c=tools).
         """
 
+        _check_duplicate_tool_names(tools)
+
         self.timeout = timeout
         self.generation_kwargs = generation_kwargs or {}
         self.url = url
         self.model = model
         self.keep_alive = keep_alive
         self.streaming_callback = streaming_callback
+        self.tools = tools
 
         self._client = Client(host=self.url, timeout=self.timeout)
 
@@ -82,6 +146,7 @@ class OllamaChatGenerator:
               Dictionary with serialized data.
         """
         callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None
+        serialized_tools = [tool.to_dict() for tool in self.tools] if self.tools else None
         return default_to_dict(
             self,
             model=self.model,
@@ -90,6 +155,7 @@ class OllamaChatGenerator:
             generation_kwargs=self.generation_kwargs,
             timeout=self.timeout,
             streaming_callback=callback_name,
+            tools=serialized_tools,
         )
 
     @classmethod
@@ -102,34 +168,14 @@ class OllamaChatGenerator:
         :returns:
             Deserialized component.
         """
+        deserialize_tools_inplace(data["init_parameters"], key="tools")
+
         init_params = data.get("init_parameters", {})
-        serialized_callback_handler = init_params.get("streaming_callback")
-        if serialized_callback_handler:
+
+        if serialized_callback_handler := init_params.get("streaming_callback"):
             data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler)
         return default_from_dict(cls, data)
 
-    def _message_to_dict(self, message: ChatMessage) -> Dict[str, str]:
-        return {"role": message.role.value, "content": message.text}
-
-    def _build_message_from_ollama_response(self, ollama_response: ChatResponse) -> ChatMessage:
-        """
-        Converts the non-streaming response from the Ollama API to a ChatMessage.
-        """
-        response_dict = ollama_response.model_dump()
-        message = ChatMessage.from_assistant(response_dict["message"]["content"])
-        message.meta.update({key: value for key, value in response_dict.items() if key != "message"})
-        return message
-
-    def _convert_to_streaming_response(self, chunks: List[StreamingChunk]) -> Dict[str, List[Any]]:
-        """
-        Converts a list of chunks response required Haystack format.
-        """
-
-        replies = [ChatMessage.from_assistant("".join([c.content for c in chunks]))]
-        meta = {key: value for key, value in chunks[0].meta.items() if key != "message"}
-
-        return {"replies": replies, "meta": [meta]}
-
     def _build_chunk(self, chunk_response: Any) -> StreamingChunk:
         """
         Converts the response from the Ollama API to a StreamingChunk.
@@ -143,23 +189,28 @@ class OllamaChatGenerator:
         chunk_message = StreamingChunk(content, meta)
         return chunk_message
 
-    def _handle_streaming_response(self, response) -> List[StreamingChunk]:
+    def _handle_streaming_response(self, response) -> Dict[str, List[Any]]:
         """
-        Handles Streaming response cases
+        Handles streaming response and converts it to Haystack format
         """
         chunks: List[StreamingChunk] = []
         for chunk in response:
-            chunk_delta: StreamingChunk = self._build_chunk(chunk)
+            chunk_delta = self._build_chunk(chunk)
             chunks.append(chunk_delta)
             if self.streaming_callback is not None:
                 self.streaming_callback(chunk_delta)
-        return chunks
+
+        replies = [ChatMessage.from_assistant("".join([c.content for c in chunks]))]
+        meta = {key: value for key, value in chunks[0].meta.items() if key != "message"}
+
+        return {"replies": replies, "meta": [meta]}
 
     @component.output_types(replies=List[ChatMessage])
     def run(
         self,
         messages: List[ChatMessage],
         generation_kwargs: Optional[Dict[str, Any]] = None,
+        tools: Optional[List[Tool]] = None,
     ):
         """
         Runs an Ollama Model on a given chat history.
@@ -170,21 +221,35 @@ class OllamaChatGenerator:
             Optional arguments to pass to the Ollama generation endpoint, such as temperature,
             top_p, etc. See the
             [Ollama docs](https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values).
-        :param streaming_callback:
-            A callback function that will be called with each response chunk in streaming mode.
+        :param tools:
+            A list of tools for which the model can prepare calls. If set, it will override the `tools` parameter set
+            during component initialization.
         :returns: A dictionary with the following keys:
             - `replies`: The responses from the model
         """
         generation_kwargs = {**self.generation_kwargs, **(generation_kwargs or {})}
 
         stream = self.streaming_callback is not None
-        messages = [self._message_to_dict(message) for message in messages]
+        tools = tools or self.tools
+        _check_duplicate_tool_names(tools)
+
+        if stream and tools:
+            msg = "Ollama does not support tools and streaming at the same time. Please choose one."
+            raise ValueError(msg)
+
+        ollama_tools = [{"type": "function", "function": {**t.tool_spec}} for t in tools] if tools else None
+
+        ollama_messages = [_convert_chatmessage_to_ollama_format(msg) for msg in messages]
         response = self._client.chat(
-            model=self.model, messages=messages, stream=stream, keep_alive=self.keep_alive, options=generation_kwargs
+            model=self.model,
+            messages=ollama_messages,
+            tools=ollama_tools,
+            stream=stream,
+            keep_alive=self.keep_alive,
+            options=generation_kwargs,
         )
 
         if stream:
-            chunks: List[StreamingChunk] = self._handle_streaming_response(response)
-            return self._convert_to_streaming_response(chunks)
+            return self._handle_streaming_response(response)
 
-        return {"replies": [self._build_message_from_ollama_response(response)]}
+        return {"replies": [_convert_ollama_response_to_chatmessage(response)]}

ollama_haystack-2.2.0/tests/test_chat_generator.py

@@ -0,0 +1,458 @@
+import json
+from unittest.mock import Mock, patch
+
+import pytest
+from haystack.components.generators.utils import print_streaming_chunk
+from haystack.dataclasses import (
+    ChatMessage,
+    ChatRole,
+    StreamingChunk,
+    TextContent,
+    ToolCall,
+)
+from haystack.tools import Tool
+from ollama._types import ChatResponse, ResponseError
+
+from haystack_integrations.components.generators.ollama.chat.chat_generator import (
+    OllamaChatGenerator,
+    _convert_chatmessage_to_ollama_format,
+    _convert_ollama_response_to_chatmessage,
+)
+
+
+@pytest.fixture
+def tools():
+    tool_parameters = {
+        "type": "object",
+        "properties": {"city": {"type": "string"}},
+        "required": ["city"],
+    }
+    tool = Tool(
+        name="weather",
+        description="useful to determine the weather in a given location",
+        parameters=tool_parameters,
+        function=lambda x: x,
+    )
+
+    return [tool]
+
+
+def test_convert_chatmessage_to_ollama_format():
+    message = ChatMessage.from_system("You are good assistant")
+    assert _convert_chatmessage_to_ollama_format(message) == {
+        "role": "system",
+        "content": "You are good assistant",
+    }
+
+    message = ChatMessage.from_user("I have a question")
+    assert _convert_chatmessage_to_ollama_format(message) == {
+        "role": "user",
+        "content": "I have a question",
+    }
+
+    message = ChatMessage.from_assistant(text="I have an answer", meta={"finish_reason": "stop"})
+    assert _convert_chatmessage_to_ollama_format(message) == {
+        "role": "assistant",
+        "content": "I have an answer",
+    }
+
+    message = ChatMessage.from_assistant(
+        tool_calls=[ToolCall(id="123", tool_name="weather", arguments={"city": "Paris"})]
+    )
+    assert _convert_chatmessage_to_ollama_format(message) == {
+        "role": "assistant",
+        "tool_calls": [
+            {
+                "type": "function",
+                "function": {"name": "weather", "arguments": {"city": "Paris"}},
+            }
+        ],
+    }
+
+    tool_result = json.dumps({"weather": "sunny", "temperature": "25"})
+    message = ChatMessage.from_tool(
+        tool_result=tool_result,
+        origin=ToolCall(tool_name="weather", arguments={"city": "Paris"}),
+    )
+    assert _convert_chatmessage_to_ollama_format(message) == {
+        "role": "tool",
+        "content": tool_result,
+    }
+
+
+def test_convert_chatmessage_to_ollama_invalid():
+    message = ChatMessage(_role=ChatRole.ASSISTANT, _content=[])
+    with pytest.raises(ValueError):
+        _convert_chatmessage_to_ollama_format(message)
+
+    message = ChatMessage(
+        _role=ChatRole.ASSISTANT,
+        _content=[
+            TextContent(text="I have an answer"),
+            TextContent(text="I have another answer"),
+        ],
+    )
+    with pytest.raises(ValueError):
+        _convert_chatmessage_to_ollama_format(message)
+
+
+def test_convert_ollama_response_to_chatmessage():
+    model = "some_model"
+
+    ollama_response = ChatResponse(
+        model=model,
+        created_at="2023-12-12T14:13:43.416799Z",
+        message={"role": "assistant", "content": "Hello! How are you today?"},
+        done=True,
+        total_duration=5191566416,
+        load_duration=2154458,
+        prompt_eval_count=26,
+        prompt_eval_duration=383809000,
+        eval_count=298,
+        eval_duration=4799921000,
+    )
+
+    observed = _convert_ollama_response_to_chatmessage(ollama_response)
+
+    assert observed.role == "assistant"
+    assert observed.text == "Hello! How are you today?"
+
+
+def test_convert_ollama_response_to_chatmessage_with_tools():
+    model = "some_model"
+
+    ollama_response = ChatResponse(
+        model=model,
+        created_at="2023-12-12T14:13:43.416799Z",
+        message={
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "function": {
+                        "name": "get_current_weather",
+                        "arguments": {"format": "celsius", "location": "Paris, FR"},
+                    }
+                }
+            ],
+        },
+        done=True,
+        total_duration=5191566416,
+        load_duration=2154458,
+        prompt_eval_count=26,
+        prompt_eval_duration=383809000,
+        eval_count=298,
+        eval_duration=4799921000,
+    )
+
+    observed = _convert_ollama_response_to_chatmessage(ollama_response)
+
+    assert observed.role == "assistant"
+    assert observed.text == ""
+    assert observed.tool_call == ToolCall(
+        tool_name="get_current_weather",
+        arguments={"format": "celsius", "location": "Paris, FR"},
+    )
+
+
+class TestOllamaChatGenerator:
+    def test_init_default(self):
+        component = OllamaChatGenerator()
+        assert component.model == "orca-mini"
+        assert component.url == "http://localhost:11434"
+        assert component.generation_kwargs == {}
+        assert component.timeout == 120
+        assert component.streaming_callback is None
+        assert component.tools is None
+        assert component.keep_alive is None
+
+    def test_init(self, tools):
+        component = OllamaChatGenerator(
+            model="llama2",
+            url="http://my-custom-endpoint:11434",
+            generation_kwargs={"temperature": 0.5},
+            timeout=5,
+            keep_alive="10m",
+            streaming_callback=print_streaming_chunk,
+            tools=tools,
+        )
+
+        assert component.model == "llama2"
+        assert component.url == "http://my-custom-endpoint:11434"
+        assert component.generation_kwargs == {"temperature": 0.5}
+        assert component.timeout == 5
+        assert component.keep_alive == "10m"
+        assert component.streaming_callback is print_streaming_chunk
+        assert component.tools == tools
+
+    def test_init_fail_with_duplicate_tool_names(self, tools):
+
+        duplicate_tools = [tools[0], tools[0]]
+        with pytest.raises(ValueError):
+            OllamaChatGenerator(tools=duplicate_tools)
+
+    def test_to_dict(self):
+        tool = Tool(
+            name="name",
+            description="description",
+            parameters={"x": {"type": "string"}},
+            function=print,
+        )
+
+        component = OllamaChatGenerator(
+            model="llama2",
+            streaming_callback=print_streaming_chunk,
+            url="custom_url",
+            generation_kwargs={"max_tokens": 10, "some_test_param": "test-params"},
+            tools=[tool],
+            keep_alive="5m",
+        )
+        data = component.to_dict()
+        assert data == {
+            "type": "haystack_integrations.components.generators.ollama.chat.chat_generator.OllamaChatGenerator",
+            "init_parameters": {
+                "timeout": 120,
+                "model": "llama2",
+                "url": "custom_url",
+                "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
+                "keep_alive": "5m",
+                "generation_kwargs": {
+                    "max_tokens": 10,
+                    "some_test_param": "test-params",
+                },
+                "tools": [
+                    {
+                        "type": "haystack.tools.tool.Tool",
+                        "data": {
+                            "description": "description",
+                            "function": "builtins.print",
+                            "name": "name",
+                            "parameters": {
+                                "x": {
+                                    "type": "string",
+                                },
+                            },
+                        },
+                    },
+                ],
+            },
+        }
+
+    def test_from_dict(self):
+        tool = Tool(
+            name="name",
+            description="description",
+            parameters={"x": {"type": "string"}},
+            function=print,
+        )
+
+        data = {
+            "type": "haystack_integrations.components.generators.ollama.chat.chat_generator.OllamaChatGenerator",
+            "init_parameters": {
+                "timeout": 120,
+                "model": "llama2",
+                "url": "custom_url",
+                "keep_alive": "5m",
+                "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
+                "generation_kwargs": {
+                    "max_tokens": 10,
+                    "some_test_param": "test-params",
+                },
+                "tools": [
+                    {
+                        "type": "haystack.tools.tool.Tool",
+                        "data": {
+                            "description": "description",
+                            "function": "builtins.print",
+                            "name": "name",
+                            "parameters": {
+                                "x": {
+                                    "type": "string",
+                                },
+                            },
+                        },
+                    },
+                ],
+            },
+        }
+        component = OllamaChatGenerator.from_dict(data)
+        assert component.model == "llama2"
+        assert component.streaming_callback is print_streaming_chunk
+        assert component.url == "custom_url"
+        assert component.keep_alive == "5m"
+        assert component.generation_kwargs == {
+            "max_tokens": 10,
+            "some_test_param": "test-params",
+        }
+        assert component.timeout == 120
+        assert component.tools == [tool]
+
+    @patch("haystack_integrations.components.generators.ollama.chat.chat_generator.Client")
+    def test_run(self, mock_client):
+        generator = OllamaChatGenerator()
+
+        mock_response = ChatResponse(
+            model="llama3.2",
+            created_at="2023-12-12T14:13:43.416799Z",
+            message={
+                "role": "assistant",
+                "content": "Fine. How can I help you today?",
+            },
+            done=True,
+            total_duration=5191566416,
+            load_duration=2154458,
+            prompt_eval_count=26,
+            prompt_eval_duration=383809000,
+            eval_count=298,
+            eval_duration=4799921000,
+        )
+
+        mock_client_instance = mock_client.return_value
+        mock_client_instance.chat.return_value = mock_response
+
+        result = generator.run(messages=[ChatMessage.from_user("Hello! How are you today?")])
+
+        mock_client_instance.chat.assert_called_once_with(
+            model="orca-mini",
+            messages=[{"role": "user", "content": "Hello! How are you today?"}],
+            stream=False,
+            tools=None,
+            options={},
+            keep_alive=None,
+        )
+
+        assert "replies" in result
+        assert len(result["replies"]) == 1
+        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_streaming(self, mock_client):
+        streaming_callback_called = False
+
+        def streaming_callback(_: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+
+        generator = OllamaChatGenerator(streaming_callback=streaming_callback)
+
+        mock_response = iter(
+            [
+                ChatResponse(
+                    model="llama3.2",
+                    created_at="2023-12-12T14:13:43.416799Z",
+                    message={"role": "assistant", "content": "first chunk "},
+                    done=False,
+                ),
+                ChatResponse(
+                    model="llama3.2",
+                    created_at="2023-12-12T14:13:43.416799Z",
+                    message={"role": "assistant", "content": "second chunk"},
+                    done=True,
+                    total_duration=4883583458,
+                    load_duration=1334875,
+                    prompt_eval_count=26,
+                    prompt_eval_duration=342546000,
+                    eval_count=282,
+                    eval_duration=4535599000,
+                ),
+            ]
+        )
+
+        mock_client_instance = mock_client.return_value
+        mock_client_instance.chat.return_value = mock_response
+
+        result = generator.run(messages=[ChatMessage.from_user("irrelevant")])
+
+        assert streaming_callback_called
+
+        assert "replies" in result
+        assert len(result["replies"]) == 1
+        assert result["replies"][0].text == "first chunk second chunk"
+        assert result["replies"][0].role == "assistant"
+
+    def test_run_fail_with_tools_and_streaming(self, tools):
+        component = OllamaChatGenerator(tools=tools, streaming_callback=print_streaming_chunk)
+
+        with pytest.raises(ValueError):
+            message = ChatMessage.from_user("irrelevant")
+            component.run([message])
+
+    @pytest.mark.integration
+    def test_live_run(self):
+        chat_generator = OllamaChatGenerator(model="llama3.2:3b")
+
+        user_questions_and_assistant_answers = [
+            ("What's the capital of France?", "Paris"),
+            ("What is the capital of Canada?", "Ottawa"),
+            ("What is the capital of Ghana?", "Accra"),
+        ]
+
+        for question, answer in user_questions_and_assistant_answers:
+            message = ChatMessage.from_user(question)
+
+            response = chat_generator.run([message])
+
+            assert isinstance(response, dict)
+            assert isinstance(response["replies"], list)
+            assert answer in response["replies"][0].text
+
+    @pytest.mark.integration
+    def test_run_with_chat_history(self):
+        chat_generator = OllamaChatGenerator(model="llama3.2:3b")
+
+        chat_messages = [
+            ChatMessage.from_user("What is the largest city in the United Kingdom by population?"),
+            ChatMessage.from_assistant("London is the largest city in the United Kingdom by population"),
+            ChatMessage.from_user("And what is the second largest?"),
+        ]
+
+        response = chat_generator.run(chat_messages)
+
+        assert isinstance(response, dict)
+        assert isinstance(response["replies"], list)
+
+        assert any(city in response["replies"][-1].text for city in ["Manchester", "Birmingham", "Glasgow"])
+
+    @pytest.mark.integration
+    def test_run_model_unavailable(self):
+        component = OllamaChatGenerator(model="unknown_model")
+
+        with pytest.raises(ResponseError):
+            message = ChatMessage.from_user("irrelevant")
+            component.run([message])
+
+    @pytest.mark.integration
+    def test_run_with_streaming(self):
+        streaming_callback = Mock()
+        chat_generator = OllamaChatGenerator(model="llama3.2:3b", streaming_callback=streaming_callback)
+
+        chat_messages = [
+            ChatMessage.from_user("What is the largest city in the United Kingdom by population?"),
+            ChatMessage.from_assistant("London is the largest city in the United Kingdom by population"),
+            ChatMessage.from_user("And what is the second largest?"),
+        ]
+
+        response = chat_generator.run(chat_messages)
+
+        streaming_callback.assert_called()
+
+        assert isinstance(response, dict)
+        assert isinstance(response["replies"], list)
+        assert any(city in response["replies"][-1].text for city in ["Manchester", "Birmingham", "Glasgow"])
+
+    @pytest.mark.integration
+    def test_run_with_tools(self, tools):
+        chat_generator = OllamaChatGenerator(model="llama3.2:3b", tools=tools)
+
+        message = ChatMessage.from_user("What is the weather in Paris?")
+        response = chat_generator.run([message])
+
+        assert len(response["replies"]) == 1
+        message = response["replies"][0]
+
+        assert message.tool_calls
+        tool_call = message.tool_call
+        assert isinstance(tool_call, ToolCall)
+        assert tool_call.tool_name == "weather"
+        assert tool_call.arguments == {"city": "Paris"}

{ollama_haystack-2.1.2 → ollama_haystack-2.2.0}/tests/test_generator.py

@@ -19,7 +19,7 @@ class TestOllamaGenerator:
             ("What is the capital of Ghana?", "Accra"),
         ]
 
-        component = OllamaGenerator()
+        component = OllamaGenerator(model="llama3.2:3b")
 
         for prompt, answer in prompts_and_answers:
             results = component.run(prompt=prompt)
@@ -147,7 +147,7 @@ class TestOllamaGenerator:
                 return chunk
 
         callback = Callback()
-        component = OllamaGenerator(streaming_callback=callback)
+        component = OllamaGenerator(model="llama3.2:3b", streaming_callback=callback)
         results = component.run(prompt="What's the capital of Netherlands?")
 
         assert len(results["replies"]) == 1

ollama_haystack-2.1.2/tests/test_chat_generator.py

@@ -1,169 +0,0 @@
-from typing import List
-from unittest.mock import Mock
-
-import pytest
-from haystack.components.generators.utils import print_streaming_chunk
-from haystack.dataclasses import ChatMessage
-from ollama._types import ChatResponse, ResponseError
-
-from haystack_integrations.components.generators.ollama import OllamaChatGenerator
-
-
-@pytest.fixture
-def chat_messages() -> List[ChatMessage]:
-    return [
-        ChatMessage.from_user("Tell me about why Super Mario is the greatest superhero"),
-        ChatMessage.from_assistant(
-            "Super Mario has prevented Bowser from destroying the world", {"something": "something"}
-        ),
-    ]
-
-
-class TestOllamaChatGenerator:
-    def test_init_default(self):
-        component = OllamaChatGenerator()
-        assert component.model == "orca-mini"
-        assert component.url == "http://localhost:11434"
-        assert component.generation_kwargs == {}
-        assert component.timeout == 120
-        assert component.keep_alive is None
-
-    def test_init(self):
-        component = OllamaChatGenerator(
-            model="llama2",
-            url="http://my-custom-endpoint:11434",
-            generation_kwargs={"temperature": 0.5},
-            keep_alive="10m",
-            timeout=5,
-        )
-
-        assert component.model == "llama2"
-        assert component.url == "http://my-custom-endpoint:11434"
-        assert component.generation_kwargs == {"temperature": 0.5}
-        assert component.timeout == 5
-        assert component.keep_alive == "10m"
-
-    def test_to_dict(self):
-        component = OllamaChatGenerator(
-            model="llama2",
-            streaming_callback=print_streaming_chunk,
-            url="custom_url",
-            generation_kwargs={"max_tokens": 10, "some_test_param": "test-params"},
-            keep_alive="5m",
-        )
-        data = component.to_dict()
-        assert data == {
-            "type": "haystack_integrations.components.generators.ollama.chat.chat_generator.OllamaChatGenerator",
-            "init_parameters": {
-                "timeout": 120,
-                "model": "llama2",
-                "keep_alive": "5m",
-                "url": "custom_url",
-                "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
-                "generation_kwargs": {"max_tokens": 10, "some_test_param": "test-params"},
-            },
-        }
-
-    def test_from_dict(self):
-        data = {
-            "type": "haystack_integrations.components.generators.ollama.chat.chat_generator.OllamaChatGenerator",
-            "init_parameters": {
-                "timeout": 120,
-                "model": "llama2",
-                "url": "custom_url",
-                "keep_alive": "5m",
-                "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",
-                "generation_kwargs": {"max_tokens": 10, "some_test_param": "test-params"},
-            },
-        }
-        component = OllamaChatGenerator.from_dict(data)
-        assert component.model == "llama2"
-        assert component.streaming_callback is print_streaming_chunk
-        assert component.url == "custom_url"
-        assert component.generation_kwargs == {"max_tokens": 10, "some_test_param": "test-params"}
-        assert component.keep_alive == "5m"
-
-    def test_build_message_from_ollama_response(self):
-        model = "some_model"
-
-        ollama_response = ChatResponse(
-            model=model,
-            created_at="2023-12-12T14:13:43.416799Z",
-            message={"role": "assistant", "content": "Hello! How are you today?"},
-            done=True,
-            total_duration=5191566416,
-            load_duration=2154458,
-            prompt_eval_count=26,
-            prompt_eval_duration=383809000,
-            eval_count=298,
-            eval_duration=4799921000,
-        )
-
-        observed = OllamaChatGenerator(model=model)._build_message_from_ollama_response(ollama_response)
-
-        assert observed.role == "assistant"
-        assert observed.text == "Hello! How are you today?"
-
-    @pytest.mark.integration
-    def test_run(self):
-        chat_generator = OllamaChatGenerator()
-
-        user_questions_and_assistant_answers = [
-            ("What's the capital of France?", "Paris"),
-            ("What is the capital of Canada?", "Ottawa"),
-            ("What is the capital of Ghana?", "Accra"),
-        ]
-
-        for question, answer in user_questions_and_assistant_answers:
-            message = ChatMessage.from_user(question)
-
-            response = chat_generator.run([message])
-
-            assert isinstance(response, dict)
-            assert isinstance(response["replies"], list)
-            assert answer in response["replies"][0].text
-
-    @pytest.mark.integration
-    def test_run_with_chat_history(self):
-        chat_generator = OllamaChatGenerator()
-
-        chat_history = [
-            ChatMessage.from_user("What is the largest city in the United Kingdom by population?"),
-            ChatMessage.from_assistant("London is the largest city in the United Kingdom by population"),
-            ChatMessage.from_user("And what is the second largest?"),
-        ]
-
-        response = chat_generator.run(chat_history)
-
-        assert isinstance(response, dict)
-        assert isinstance(response["replies"], list)
-        assert "Manchester" in response["replies"][-1].text or "Glasgow" in response["replies"][-1].text
-
-    @pytest.mark.integration
-    def test_run_model_unavailable(self):
-        component = OllamaChatGenerator(model="Alistair_and_Stefano_are_great")
-
-        with pytest.raises(ResponseError):
-            message = ChatMessage.from_user(
-                "Based on your infinite wisdom, can you tell me why Alistair and Stefano are so great?"
-            )
-            component.run([message])
-
-    @pytest.mark.integration
-    def test_run_with_streaming(self):
-        streaming_callback = Mock()
-        chat_generator = OllamaChatGenerator(streaming_callback=streaming_callback)
-
-        chat_history = [
-            ChatMessage.from_user("What is the largest city in the United Kingdom by population?"),
-            ChatMessage.from_assistant("London is the largest city in the United Kingdom by population"),
-            ChatMessage.from_user("And what is the second largest?"),
-        ]
-
-        response = chat_generator.run(chat_history)
-
-        streaming_callback.assert_called()
-
-        assert isinstance(response, dict)
-        assert isinstance(response["replies"], list)
-        assert "Manchester" in response["replies"][-1].text or "Glasgow" in response["replies"][-1].text