feat: add to_dict and from_dict to StreamingChunk (#9608)

* add to_dict and from_dict to StreamingChunk

* bugfix: set index field

* release notes + minor bugfix in tests

* ensure full serialization + update tests

* support de/serialization for all involved dataclasses

* remove unnecessary import

* code cleanup + only allow dicts in from_dict

* update release notes

* add more tests

Author: srishti-git1110

haystack/dataclasses/chat_message.py

@@ -60,6 +60,26 @@ class ToolCall:
     arguments: Dict[str, Any]
     id: Optional[str] = None  # noqa: A003
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert ToolCall into a dictionary.
+
+        :returns: A dictionary with keys 'tool_name', 'arguments', and 'id'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ToolCall":
+        """
+        Creates a new ToolCall object from a dictionary.
+
+        :param data:
+            The dictionary to build the ToolCall object.
+        :returns:
+            The created object.
+        """
+        return ToolCall(**data)
+
 
 @dataclass
 class ToolCallResult:
@@ -75,6 +95,31 @@ class ToolCallResult:
     origin: ToolCall
     error: bool
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Converts ToolCallResult into a dictionary.
+
+        :returns: A dictionary with keys 'result', 'origin', and 'error'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ToolCallResult":
+        """
+        Creates a ToolCallResult from a dictionary.
+
+        :param data:
+            The dictionary to build the ToolCallResult object.
+        :returns:
+            The created object.
+        """
+        if not all(x in data for x in ["result", "origin", "error"]):
+            raise ValueError(
+                "Fields `result`, `origin`, `error` are required for ToolCallResult deserialization. "
+                f"Received dictionary with keys {list(data.keys())}"
+            )
+        return ToolCallResult(result=data["result"], origin=ToolCall.from_dict(data["origin"]), error=data["error"])
+
 
 @dataclass
 class TextContent:

haystack/dataclasses/streaming_chunk.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from dataclasses import dataclass, field
+from dataclasses import asdict, dataclass, field
 from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Union, overload
 
 from haystack.core.component import Component
@@ -30,6 +30,24 @@ class ToolCallDelta:
     arguments: Optional[str] = field(default=None)
     id: Optional[str] = field(default=None)  # noqa: A003
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Returns a dictionary representation of the ToolCallDelta.
+
+        :returns: A dictionary with keys 'index', 'tool_name', 'arguments', and 'id'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ToolCallDelta":
+        """
+        Creates a ToolCallDelta from a serialized representation.
+
+        :param data: Dictionary containing ToolCallDelta's attributes.
+        :returns: A ToolCallDelta instance.
+        """
+        return ToolCallDelta(**data)
+
 
 @dataclass
 class ComponentInfo:
@@ -58,6 +76,24 @@ def from_component(cls, component: Component) -> "ComponentInfo":
         component_name = getattr(component, "__component_name__", None)
         return cls(type=component_type, name=component_name)
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Returns a dictionary representation of ComponentInfo.
+
+        :returns: A dictionary with keys 'type' and 'name'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ComponentInfo":
+        """
+        Creates a ComponentInfo from a serialized representation.
+
+        :param data: Dictionary containing ComponentInfo's attributes.
+        :returns: A ComponentInfo instance.
+        """
+        return ComponentInfo(**data)
+
 
 @dataclass
 class StreamingChunk:
@@ -102,6 +138,47 @@ def __post_init__(self):
         if (self.tool_calls or self.tool_call_result) and self.index is None:
             raise ValueError("If `tool_call`, or `tool_call_result` is set, `index` must also be set.")
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Returns a dictionary representation of the StreamingChunk.
+
+        :returns: Serialized dictionary representation of the calling object.
+        """
+        return {
+            "content": self.content,
+            "meta": self.meta,
+            "component_info": self.component_info.to_dict() if self.component_info else None,
+            "index": self.index,
+            "tool_calls": [tc.to_dict() for tc in self.tool_calls] if self.tool_calls else None,
+            "tool_call_result": self.tool_call_result.to_dict() if self.tool_call_result else None,
+            "start": self.start,
+            "finish_reason": self.finish_reason,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "StreamingChunk":
+        """
+        Creates a deserialized StreamingChunk instance from a serialized representation.
+
+        :param data: Dictionary containing the StreamingChunk's attributes.
+        :returns: A StreamingChunk instance.
+        """
+        if "content" not in data:
+            raise ValueError("Missing required field `content` in StreamingChunk deserialization.")
+
+        return StreamingChunk(
+            content=data["content"],
+            meta=data.get("meta", {}),
+            component_info=ComponentInfo.from_dict(data["component_info"]) if data.get("component_info") else None,
+            index=data.get("index"),
+            tool_calls=[ToolCallDelta.from_dict(tc) for tc in data["tool_calls"]] if data.get("tool_calls") else None,
+            tool_call_result=ToolCallResult.from_dict(data["tool_call_result"])
+            if data.get("tool_call_result")
+            else None,
+            start=data.get("start", False),
+            finish_reason=data.get("finish_reason"),
+        )
+
 
 SyncStreamingCallbackT = Callable[[StreamingChunk], None]
 AsyncStreamingCallbackT = Callable[[StreamingChunk], Awaitable[None]]

releasenotes/notes/add-to-dict-from-dict-streamingchunk-f42dfbcca72e0465.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Add `to_dict` and `from_dict` to classes StreamingChunk, ToolCallResult, ToolCall, ComponentInfo, and ToolCallDelta to make it consistent with our other dataclasses in having serialization and deserialization methods.

test/dataclasses/test_streaming_chunk.py

@@ -135,3 +135,121 @@ def test_finish_reason_tool_call_results():
     assert chunk.finish_reason == "tool_call_results"
     assert chunk.meta["finish_reason"] == "tool_call_results"
     assert chunk.content == ""
+
+
+def test_to_dict_tool_call_result():
+    """Test the to_dict method for StreamingChunk with tool_call_result."""
+    component = ExampleComponent()
+    component_info = ComponentInfo.from_component(component)
+    tool_call_result = ToolCallResult(
+        result="output", origin=ToolCall(id="123", tool_name="test_tool", arguments={"arg1": "value1"}), error=False
+    )
+
+    chunk = StreamingChunk(
+        content="",
+        meta={"key": "value"},
+        index=0,
+        component_info=component_info,
+        tool_call_result=tool_call_result,
+        finish_reason="tool_call_results",
+    )
+
+    d = chunk.to_dict()
+
+    assert d["content"] == ""
+    assert d["meta"] == {"key": "value"}
+    assert d["index"] == 0
+    assert d["component_info"]["type"] == "test_streaming_chunk.ExampleComponent"
+    assert d["tool_call_result"]["result"] == "output"
+    assert d["tool_call_result"]["error"] is False
+    assert d["tool_call_result"]["origin"]["id"] == "123"
+    assert d["tool_call_result"]["origin"]["arguments"]["arg1"] == "value1"
+    assert d["finish_reason"] == "tool_call_results"
+
+
+def test_to_dict_tool_calls():
+    """Test the to_dict method for StreamingChunk with tool_calls."""
+    component = ExampleComponent()
+    component_info = ComponentInfo.from_component(component)
+    tool_calls = [
+        ToolCallDelta(id="123", tool_name="test_tool", arguments='{"arg1": "value1"}', index=0),
+        ToolCallDelta(id="456", tool_name="another_tool", arguments='{"arg2": "value2"}', index=1),
+    ]
+
+    chunk = StreamingChunk(
+        content="",
+        meta={"key": "value"},
+        index=0,
+        component_info=component_info,
+        tool_calls=tool_calls,
+        finish_reason="tool_calls",
+    )
+
+    d = chunk.to_dict()
+
+    assert d["content"] == ""
+    assert d["meta"] == {"key": "value"}
+    assert d["index"] == 0
+    assert d["component_info"]["type"] == "test_streaming_chunk.ExampleComponent"
+    assert len(d["tool_calls"]) == 2
+    assert d["tool_calls"][0]["id"] == "123"
+    assert d["tool_calls"][0]["index"] == 0
+    assert d["tool_calls"][1]["id"] == "456"
+    assert d["tool_calls"][1]["index"] == 1
+    assert d["finish_reason"] == "tool_calls"
+
+
+def test_from_dict_tool_call_result():
+    """Test the from_dict method for StreamingChunk with tool_call_result."""
+    component_info = {"type": "test_streaming_chunk.ExampleComponent", "name": "test_component"}
+    tool_call_result = {
+        "result": "output",
+        "origin": {"id": "123", "tool_name": "test_tool", "arguments": {"arg1": "value1"}},
+        "error": False,
+    }
+
+    data = {
+        "content": "",
+        "meta": {"key": "value"},
+        "index": 0,
+        "component_info": component_info,
+        "tool_call_result": tool_call_result,
+        "finish_reason": "tool_call_results",
+    }
+
+    chunk = StreamingChunk.from_dict(data)
+
+    assert chunk.content == ""
+    assert chunk.meta == {"key": "value"}
+    assert chunk.index == 0
+    assert chunk.component_info.type == "test_streaming_chunk.ExampleComponent"
+    assert chunk.component_info.name == "test_component"
+    assert chunk.tool_call_result.result == "output"
+    assert chunk.tool_call_result.error is False
+    assert chunk.tool_call_result.origin.id == "123"
+
+
+def test_from_dict_tool_calls():
+    """Test the from_dict method for StreamingChunk with tool_calls."""
+    component_info = {"type": "test_streaming_chunk.ExampleComponent", "name": "test_component"}
+    tool_calls = [{"id": "123", "tool_name": "test_tool", "arguments": '{"arg1": "value1"}', "index": 0}]
+
+    data = {
+        "content": "",
+        "meta": {"key": "value"},
+        "index": 0,
+        "component_info": component_info,
+        "tool_calls": tool_calls,
+        "finish_reason": "tool_calls",
+    }
+
+    chunk = StreamingChunk.from_dict(data)
+
+    assert chunk.content == ""
+    assert chunk.meta == {"key": "value"}
+    assert chunk.index == 0
+    assert chunk.component_info.type == "test_streaming_chunk.ExampleComponent"
+    assert chunk.component_info.name == "test_component"
+    assert chunk.tool_calls[0].tool_name == "test_tool"
+    assert chunk.tool_calls[0].index == 0
+    assert chunk.finish_reason == "tool_calls"