feat: pass component_info to StreamingChunk in OllamaChatGenerator (#2039)

* feat: pass component_info to StreamingChunk in OllamaChatGenerator

* small tests improvements

* pin to patch release

Author: anakin87

integrations/ollama/pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.13.1", "ollama>=0.4.0", "pydantic"]
+dependencies = ["haystack-ai>=2.15.1", "ollama>=0.4.0", "pydantic"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/ollama#readme"
@@ -158,9 +158,7 @@ show_missing = true
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.pytest.ini_options]
-markers = [
-    "integration: marks tests as slow (deselect with '-m \"not integration\"')",
-]
+markers = ["integration: integration tests"]
 log_cli = true
 addopts = ["--import-mode=importlib"]
 asyncio_mode = "auto"

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

@@ -1,7 +1,7 @@
-from typing import Any, Callable, Dict, List, Literal, Optional, Union
+from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Union
 
 from haystack import component, default_from_dict, default_to_dict
-from haystack.dataclasses import ChatMessage, StreamingChunk, ToolCall
+from haystack.dataclasses import ChatMessage, ComponentInfo, StreamingChunk, ToolCall
 from haystack.tools import (
     Tool,
     _check_duplicate_tool_names,
@@ -100,7 +100,7 @@ def _convert_ollama_meta_to_openai_format(input_response_dict: Dict) -> Dict:
     return meta
 
 
-def _convert_ollama_response_to_chatmessage(ollama_response: "ChatResponse") -> ChatMessage:
+def _convert_ollama_response_to_chatmessage(ollama_response: ChatResponse) -> ChatMessage:
     """
     Convert non-streaming Ollama Chat API response to Haystack ChatMessage with the assistant role.
     """
@@ -271,7 +271,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "OllamaChatGenerator":
         return default_from_dict(cls, data)
 
     @staticmethod
-    def _build_chunk(chunk_response: Any) -> StreamingChunk:
+    def _build_chunk(chunk_response: ChatResponse, component_info: ComponentInfo) -> StreamingChunk:
         """
         Convert one Ollama stream-chunk to Haystack StreamingChunk.
         """
@@ -283,11 +283,11 @@ def _build_chunk(chunk_response: Any) -> StreamingChunk:
         if tool_calls := chunk_response_dict["message"].get("tool_calls"):
             meta["tool_calls"] = tool_calls
 
-        return StreamingChunk(content, meta)
+        return StreamingChunk(content=content, meta=meta, component_info=component_info)
 
     def _handle_streaming_response(
         self,
-        response_iter: Any,
+        response_iter: Iterator[ChatResponse],
         callback: Optional[Callable[[StreamingChunk], None]],
     ) -> Dict[str, List[ChatMessage]]:
         """
@@ -296,6 +296,8 @@ def _handle_streaming_response(
         or as full JSON dicts.
         """
 
+        component_info = ComponentInfo.from_component(self)
+
         chunks: List[StreamingChunk] = []
 
         # Accumulators
@@ -305,7 +307,7 @@ def _handle_streaming_response(
 
         # Stream
         for raw in response_iter:
-            chunk = self._build_chunk(raw)
+            chunk = self._build_chunk(chunk_response=raw, component_info=component_info)
             chunks.append(chunk)
 
             if callback:
@@ -428,8 +430,8 @@ def run(
             format=self.response_format,
         )
 
-        if is_stream:
-            return self._handle_streaming_response(response, callback)
+        if isinstance(response, Iterator):
+            return self._handle_streaming_response(response_iter=response, callback=callback)
 
         # non-stream path
-        return {"replies": [_convert_ollama_response_to_chatmessage(response)]}
+        return {"replies": [_convert_ollama_response_to_chatmessage(ollama_response=response)]}

integrations/ollama/tests/test_chat_generator.py

@@ -6,6 +6,7 @@
 from haystack.dataclasses import (
     ChatMessage,
     ChatRole,
+    ComponentInfo,
     StreamingChunk,
     TextContent,
     ToolCall,
@@ -287,6 +288,9 @@ def test_to_dict(self):
                                     "type": "string",
                                 },
                             },
+                            "outputs_to_string": None,
+                            "inputs_from_state": None,
+                            "outputs_to_state": None,
                         },
                     },
                 ],
@@ -297,15 +301,6 @@ def test_to_dict(self):
             },
         }
 
-        # add outputs_to_string, inputs_from_state and outputs_to_state tool parameters for compatibility with
-        # haystack-ai>=2.12.0
-        if hasattr(tool, "outputs_to_string"):
-            expected_dict["init_parameters"]["tools"][0]["data"]["outputs_to_string"] = tool.outputs_to_string
-        if hasattr(tool, "inputs_from_state"):
-            expected_dict["init_parameters"]["tools"][0]["data"]["inputs_from_state"] = tool.inputs_from_state
-        if hasattr(tool, "outputs_to_state"):
-            expected_dict["init_parameters"]["tools"][0]["data"]["outputs_to_state"] = tool.outputs_to_state
-
         assert data == expected_dict
 
     def test_from_dict(self):
@@ -365,6 +360,30 @@ def test_from_dict(self):
             "properties": {"name": {"type": "string"}, "age": {"type": "number"}},
         }
 
+    def test_build_chunk(self):
+        generator = OllamaChatGenerator()
+
+        mock_chunk_response = Mock()
+        mock_chunk_response.model_dump.return_value = {
+            "message": {"role": "assistant", "content": "Hello world"},
+            "model": "llama2",
+            "created_at": "2023-12-12T14:13:43.416799Z",
+            "done": False,
+        }
+
+        component_info = ComponentInfo.from_component(generator)
+
+        chunk = generator._build_chunk(mock_chunk_response, component_info)
+
+        assert isinstance(chunk, StreamingChunk)
+        assert chunk.content == "Hello world"
+        assert chunk.component_info == component_info
+        assert chunk.meta["role"] == "assistant"
+        assert chunk.meta["model"] == "llama2"
+        assert chunk.meta["created_at"] == "2023-12-12T14:13:43.416799Z"
+        assert chunk.meta["done"] is False
+        assert "tool_calls" not in chunk.meta
+
     @patch("haystack_integrations.components.generators.ollama.chat.chat_generator.Client")
     def test_run(self, mock_client):
         generator = OllamaChatGenerator()
@@ -407,11 +426,10 @@ def test_run(self, mock_client):
 
     @patch("haystack_integrations.components.generators.ollama.chat.chat_generator.Client")
     def test_run_streaming(self, mock_client):
-        streaming_callback_called = False
+        collected_chunks = []
 
-        def streaming_callback(_: StreamingChunk) -> None:
-            nonlocal streaming_callback_called
-            streaming_callback_called = True
+        def streaming_callback(chunk: StreamingChunk) -> None:
+            collected_chunks.append(chunk)
 
         generator = OllamaChatGenerator(streaming_callback=streaming_callback)
 
@@ -443,7 +461,16 @@ def streaming_callback(_: StreamingChunk) -> None:
 
         result = generator.run(messages=[ChatMessage.from_user("irrelevant")])
 
-        assert streaming_callback_called
+        assert len(collected_chunks) == 2
+        assert collected_chunks[0].content == "first chunk "
+        assert collected_chunks[1].content == "second chunk"
+
+        for chunk in collected_chunks:
+            assert (
+                chunk.component_info.type
+                == "haystack_integrations.components.generators.ollama.chat.chat_generator.OllamaChatGenerator"
+            )
+            assert chunk.component_info.name is None  # not in a pipeline
 
         assert "replies" in result
         assert len(result["replies"]) == 1