feat: add ReasoningContent to ChatMessage (#9696)

* feat: add ReasoningContent to ChatMessage

* more tests

* release note

* Update haystack/dataclasses/chat_message.py

Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com>

---------

Co-authored-by: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com>

Author: anakin87

haystack/dataclasses/chat_message.py

@@ -145,13 +145,48 @@ def from_dict(cls, data: dict[str, Any]) -> "TextContent":
         return TextContent(**data)
 
 
-ChatMessageContentT = Union[TextContent, ToolCall, ToolCallResult, ImageContent]
+@dataclass
+class ReasoningContent:
+    """
+    Represents the optional reasoning content prepared by the model, usually contained in an assistant message.
+
+    :param reasoning_text: The reasoning text produced by the model.
+    :param extra: Dictionary of extra information about the reasoning content. Use to store provider-specific
+        information. To avoid serialization issues, values should be JSON serializable.
+    """
+
+    reasoning_text: str
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert ReasoningContent into a dictionary.
+
+        :returns: A dictionary with keys 'reasoning_text', and 'extra'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ReasoningContent":
+        """
+        Creates a new ReasoningContent object from a dictionary.
+
+        :param data:
+            The dictionary to build the ReasoningContent object.
+        :returns:
+            The created object.
+        """
+        return ReasoningContent(**data)
+
+
+ChatMessageContentT = Union[TextContent, ToolCall, ToolCallResult, ImageContent, ReasoningContent]
 
 _CONTENT_PART_CLASSES_TO_SERIALIZATION_KEYS: dict[type[ChatMessageContentT], str] = {
     TextContent: "text",
     ToolCall: "tool_call",
     ToolCallResult: "tool_call_result",
     ImageContent: "image",
+    ReasoningContent: "reasoning",
 }
 
 
@@ -200,7 +235,7 @@ def _serialize_content_part(part: ChatMessageContentT) -> dict[str, Any]:
 
 
 @dataclass
-class ChatMessage:
+class ChatMessage:  # pylint: disable=too-many-public-methods # it's OK since we expose several properties
     """
     Represents a message in a LLM chat conversation.
 
@@ -334,6 +369,22 @@ def image(self) -> Optional[ImageContent]:
             return images[0]
         return None
 
+    @property
+    def reasonings(self) -> list[ReasoningContent]:
+        """
+        Returns the list of all reasoning contents contained in the message.
+        """
+        return [content for content in self._content if isinstance(content, ReasoningContent)]
+
+    @property
+    def reasoning(self) -> Optional[ReasoningContent]:
+        """
+        Returns the first reasoning content contained in the message.
+        """
+        if reasonings := self.reasonings:
+            return reasonings[0]
+        return None
+
     def is_from(self, role: Union[ChatRole, str]) -> bool:
         """
         Check if the message is from a specific role.
@@ -406,17 +457,27 @@ def from_assistant(
         meta: Optional[dict[str, Any]] = None,
         name: Optional[str] = None,
         tool_calls: Optional[list[ToolCall]] = None,
+        *,
+        reasoning: Optional[Union[str, ReasoningContent]] = None,
     ) -> "ChatMessage":
         """
         Create a message from the assistant.
 
         :param text: The text content of the message.
         :param meta: Additional metadata associated with the message.
-        :param tool_calls: The Tool calls to include in the message.
         :param name: An optional name for the participant. This field is only supported by OpenAI.
+        :param tool_calls: The Tool calls to include in the message.
+        :param reasoning: The reasoning content to include in the message.
         :returns: A new ChatMessage instance.
         """
         content: list[ChatMessageContentT] = []
+        if reasoning:
+            if isinstance(reasoning, str):
+                content.append(ReasoningContent(reasoning_text=reasoning))
+            elif isinstance(reasoning, ReasoningContent):
+                content.append(reasoning)
+            else:
+                raise TypeError(f"reasoning must be a string or a ReasoningContent object, got {type(reasoning)}")
         if text is not None:
             content.append(TextContent(text=text))
         if tool_calls:
@@ -576,6 +637,7 @@ def to_openai_dict_format(self, require_tool_call_ids: bool = True) -> dict[str,
             return openai_msg
 
         # system and assistant messages
+        # OpenAI Chat Completions API does not support reasoning content, so we ignore it
         if text_contents:
             openai_msg["content"] = text_contents[0]
         if tool_calls:

haystack/utils/jinja2_chat_extension.py

@@ -14,6 +14,7 @@
     ChatMessageContentT,
     ChatRole,
     ImageContent,
+    ReasoningContent,
     TextContent,
     ToolCall,
     ToolCallResult,
@@ -236,14 +237,19 @@ def _validate_build_chat_message(
         if role == "assistant":
             texts = [part.text for part in parts if isinstance(part, TextContent)]
             tool_calls = [part for part in parts if isinstance(part, ToolCall)]
+            reasoning = [part for part in parts if isinstance(part, ReasoningContent)]
             if len(texts) > 1:
                 raise ValueError("Assistant message must contain one text part at most.")
             if len(texts) == 0 and len(tool_calls) == 0:
                 raise ValueError("Assistant message must contain at least one text or tool call part.")
-            if len(parts) > len(texts) + len(tool_calls):
-                raise ValueError("Assistant message must contain only text or tool call parts.")
+            if len(parts) > len(texts) + len(tool_calls) + len(reasoning):
+                raise ValueError("Assistant message must contain only text, tool call or reasoning parts.")
             return ChatMessage.from_assistant(
-                meta=meta, name=name, text=texts[0] if texts else None, tool_calls=tool_calls or None
+                meta=meta,
+                name=name,
+                text=texts[0] if texts else None,
+                tool_calls=tool_calls or None,
+                reasoning=reasoning[0] if reasoning else None,
             )
 
         if role == "tool":

releasenotes/notes/reasoning-content-ac2d33c3a76406dc.yaml

@@ -0,0 +1,9 @@
+---
+features:
+  - |
+    Added `ReasoningContent` as a new content part to the `ChatMessage` dataclass. This allows storing model
+    reasoning text and additional metadata in assistant messages.
+    Assistant messages can now include reasoning content using the `reasoning` parameter in
+    `ChatMessage.from_assistant()`.
+    We will progressively update the implementations for Chat Generators with LLMs that support reasoning to use this
+    new content part.

test/components/builders/test_chat_prompt_builder.py

@@ -12,7 +12,7 @@
 from haystack import component
 from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
 from haystack.core.pipeline.pipeline import Pipeline
-from haystack.dataclasses.chat_message import ChatMessage, ImageContent
+from haystack.dataclasses.chat_message import ChatMessage, ImageContent, ReasoningContent
 from haystack.dataclasses.document import Document
 
 
@@ -891,6 +891,26 @@ def test_run_multiple_images(self, base64_image_string):
             )
         ]
 
+    def test_run_reasoning(self):
+        template = """
+        {% message role="user" %}
+        Hello! I am {{user_name}}. How much is 2 + 2?
+        {% endmessage %}
+
+        {% message role="assistant" %}
+        {{ reasoning | templatize_part }}
+        The answer is 4.
+        {% endmessage %}
+        """
+        builder = ChatPromptBuilder(template=template)
+        reasoning = ReasoningContent(reasoning_text="Let me think about it...", extra={"key": "value"})
+        result = builder.run(user_name="John", reasoning=reasoning)
+
+        assert result["prompt"] == [
+            ChatMessage.from_user(text="Hello! I am John. How much is 2 + 2?"),
+            ChatMessage.from_assistant(reasoning=reasoning, text="The answer is 4."),
+        ]
+
     def test_to_dict(self):
         template = """
         {% message role="user" %}

test/dataclasses/test_chat_message.py

@@ -6,7 +6,14 @@
 
 import pytest
 
-from haystack.dataclasses.chat_message import ChatMessage, ChatRole, TextContent, ToolCall, ToolCallResult
+from haystack.dataclasses.chat_message import (
+    ChatMessage,
+    ChatRole,
+    ReasoningContent,
+    TextContent,
+    ToolCall,
+    ToolCallResult,
+)
 from haystack.dataclasses.image_content import ImageContent
 
 
@@ -80,9 +87,28 @@ def test_text_content_from_dict(self):
         tc = TextContent.from_dict({"text": "Hello"})
         assert tc.text == "Hello"
 
+    def test_reasoning_content_init(self):
+        rc = ReasoningContent(reasoning_text="Let me think about it...")
+
+        assert rc.reasoning_text == "Let me think about it..."
+        assert rc.extra == {}
+
+        rc = ReasoningContent(reasoning_text="Let me think about it...", extra={"key": "value"})
+        assert rc.reasoning_text == "Let me think about it..."
+        assert rc.extra == {"key": "value"}
+
+    def test_reasoning_content_to_dict(self):
+        rc = ReasoningContent(reasoning_text="Let me think about it...", extra={"key": "value"})
+        assert rc.to_dict() == {"reasoning_text": "Let me think about it...", "extra": {"key": "value"}}
+
+    def test_reasoning_content_from_dict(self):
+        rc = ReasoningContent.from_dict({"reasoning_text": "Let me think about it...", "extra": {"key": "value"}})
+        assert rc.reasoning_text == "Let me think about it..."
+        assert rc.extra == {"key": "value"}
+
 
 class TestChatMessage:
-    def test_from_assistant_with_valid_content(self):
+    def test_from_assistant_with_text(self):
         text = "Hello, how can I assist you?"
         message = ChatMessage.from_assistant(text)
 
@@ -99,6 +125,8 @@ def test_from_assistant_with_valid_content(self):
         assert not message.tool_call_result
         assert not message.images
         assert not message.image
+        assert not message.reasonings
+        assert not message.reasoning
 
     def test_from_assistant_with_tool_calls(self):
         tool_calls = [
@@ -120,6 +148,53 @@ def test_from_assistant_with_tool_calls(self):
         assert not message.tool_call_result
         assert not message.images
         assert not message.image
+        assert not message.reasoning
+        assert not message.reasonings
+
+    def test_from_assistant_with_reasoning_object(self):
+        reasoning = ReasoningContent(reasoning_text="Let me think about it...", extra={"key": "value"})
+        text = "After thinking about it, I can say that the answer is 42."
+        message = ChatMessage.from_assistant(text=text, reasoning=reasoning)
+
+        assert message.role == ChatRole.ASSISTANT
+        assert message._content == [reasoning, TextContent(text=text)]
+
+        assert message.texts == [text]
+        assert message.text == text
+        assert message.reasoning == reasoning
+        assert message.reasonings == [reasoning]
+
+        assert not message.tool_calls
+        assert not message.tool_call
+        assert not message.tool_call_results
+        assert not message.tool_call_result
+        assert not message.images
+        assert not message.image
+
+    def test_from_assistant_with_reasoning_string(self):
+        reasoning = "Let me think about it..."
+        text = "After thinking about it, I can say that the answer is 42."
+        message = ChatMessage.from_assistant(text=text, reasoning=reasoning)
+
+        expected_reasoning_content = ReasoningContent(reasoning_text=reasoning)
+        assert message.role == ChatRole.ASSISTANT
+        assert message._content == [expected_reasoning_content, TextContent(text=text)]
+
+        assert message.texts == [text]
+        assert message.text == text
+        assert message.reasoning == expected_reasoning_content
+        assert message.reasonings == [expected_reasoning_content]
+
+        assert not message.tool_calls
+        assert not message.tool_call
+        assert not message.tool_call_results
+        assert not message.tool_call_result
+        assert not message.images
+        assert not message.image
+
+    def test_from_assistant_with_invalid_reasoning(self):
+        with pytest.raises(TypeError):
+            ChatMessage.from_assistant(text="text", reasoning=123)
 
     def test_from_user_with_valid_content(self):
         text = "I have a question."
@@ -138,6 +213,8 @@ def test_from_user_with_valid_content(self):
         assert not message.tool_call_result
         assert not message.images
         assert not message.image
+        assert not message.reasonings
+        assert not message.reasoning
 
     def test_from_user_with_name(self):
         text = "I have a question."
@@ -207,6 +284,8 @@ def test_from_system_with_valid_content(self):
         assert not message.tool_call_result
         assert not message.images
         assert not message.image
+        assert not message.reasonings
+        assert not message.reasoning
 
     def test_from_tool_with_valid_content(self):
         tool_result = "Tool result"
@@ -227,6 +306,8 @@ def test_from_tool_with_valid_content(self):
         assert not message.text
         assert not message.images
         assert not message.image
+        assert not message.reasonings
+        assert not message.reasoning
 
     def test_multiple_text_segments(self):
         texts = [TextContent(text="Hello"), TextContent(text="World")]
@@ -266,10 +347,13 @@ def test_serde(self, base64_image_string):
             meta={"key": "value"},
             validation=True,
         )
+        reasoning_content = ReasoningContent(reasoning_text="Let me think about it...", extra={"key": "value"})
         meta = {"some": "info"}
 
         message = ChatMessage(
-            _role=role, _content=[text_content, tool_call, tool_call_result, image_content], _meta=meta
+            _role=role,
+            _content=[text_content, tool_call, tool_call_result, image_content, reasoning_content],
+            _meta=meta,
         )
 
         serialized_message = message.to_dict()
@@ -293,6 +377,7 @@ def test_serde(self, base64_image_string):
                         "validation": True,
                     }
                 },
+                {"reasoning": {"reasoning_text": "Let me think about it...", "extra": {"key": "value"}}},
             ],
             "role": "assistant",
             "name": None,