From bc42b3fad13eaaebade1d3d9856565c106bb8c27 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Wed, 15 Jul 2026 16:55:48 +0200 Subject: [PATCH] Update ChatMessage to use to_dict and from_dict in pydantic serialization --- haystack/dataclasses/chat_message.py | 31 +++++++++++ ...pydantic-integration-41204145db28d04d.yaml | 14 +++++ test/dataclasses/test_chat_message.py | 52 +++++++++++++++++++ 3 files changed, 97 insertions(+) create mode 100644 releasenotes/notes/chatmessage-pydantic-integration-41204145db28d04d.yaml diff --git a/haystack/dataclasses/chat_message.py b/haystack/dataclasses/chat_message.py index 39da893f665..9f5cda8d902 100644 --- a/haystack/dataclasses/chat_message.py +++ b/haystack/dataclasses/chat_message.py @@ -8,6 +8,9 @@ from enum import Enum from typing import Any +from pydantic import GetCoreSchemaHandler +from pydantic_core import CoreSchema, core_schema + from haystack import logging from haystack.dataclasses.file_content import FileContent from haystack.dataclasses.image_content import ImageContent @@ -267,6 +270,21 @@ def _serialize_content_part(part: ChatMessageContentT) -> dict[str, Any]: return {serialization_key: part.to_dict()} +def _validate_chat_message(value: Any) -> "ChatMessage": + """ + Coerce a value into a ChatMessage during Pydantic validation. + + :param value: A ChatMessage instance or a dictionary in a format accepted by `ChatMessage.from_dict`. + :returns: A ChatMessage object. + :raises ValueError: If the value is neither a ChatMessage nor a dictionary. + """ + if isinstance(value, ChatMessage): + return value + if isinstance(value, dict): + return ChatMessage.from_dict(value) + raise ValueError(f"Expected `ChatMessage` or `dict`, got `{type(value).__name__}`") + + @_warn_on_inplace_mutation @dataclass class ChatMessage: @@ -618,6 +636,19 @@ def from_dict(cls, data: dict[str, Any]) -> "ChatMessage": raise ValueError(f"Missing 'content' or '_content' in serialized ChatMessage: `{data}`") + @classmethod + def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + """ + Defines how Pydantic validates and serializes ChatMessage. + + Dictionaries are validated with `from_dict` and existing instances pass through unchanged. + Serialization uses `to_dict`, so Pydantic's `model_dump` produces the canonical format instead + of the raw dataclass fields. + """ + return core_schema.no_info_plain_validator_function( + _validate_chat_message, serialization=core_schema.plain_serializer_function_ser_schema(cls.to_dict) + ) + def to_openai_dict_format(self, require_tool_call_ids: bool = True) -> dict[str, Any]: """ Convert a ChatMessage to the dictionary format expected by OpenAI's Chat Completions API. diff --git a/releasenotes/notes/chatmessage-pydantic-integration-41204145db28d04d.yaml b/releasenotes/notes/chatmessage-pydantic-integration-41204145db28d04d.yaml new file mode 100644 index 00000000000..699dce0bf0c --- /dev/null +++ b/releasenotes/notes/chatmessage-pydantic-integration-41204145db28d04d.yaml @@ -0,0 +1,14 @@ +--- +upgrade: + - | + Pydantic's ``model_dump`` on models containing ``ChatMessage`` fields now returns the canonical + ``ChatMessage.to_dict`` format (``role``/``content`` keys) instead of the raw dataclass fields + (``_role``/``_content`` keys). ``ChatMessage.from_dict`` accepts both formats, so round-tripping + existing payloads keeps working. If you post-process dumped messages and rely on the raw dataclass + keys, update your code to the canonical format. +enhancements: + - | + ``ChatMessage`` now integrates natively with Pydantic. Pydantic models with ``ChatMessage`` fields + serialize with ``ChatMessage.to_dict`` and validate dictionaries with ``ChatMessage.from_dict``, so + ``model_dump`` produces the canonical serialization format and serialized messages round-trip + through Pydantic models without custom serializers or validators. diff --git a/test/dataclasses/test_chat_message.py b/test/dataclasses/test_chat_message.py index 52c600446ee..3223d2d530d 100644 --- a/test/dataclasses/test_chat_message.py +++ b/test/dataclasses/test_chat_message.py @@ -7,6 +7,7 @@ from collections.abc import Sequence import pytest +from pydantic import BaseModel, ValidationError from haystack.dataclasses.chat_message import ( ChatMessage, @@ -756,6 +757,57 @@ def test_to_trace_dict_with_file_content(self, base64_pdf_string): } +class ChatHistory(BaseModel): + messages: list[ChatMessage] + + +class TestChatMessagePydanticIntegration: + def test_model_dump_uses_to_dict(self): + messages = [ + ChatMessage.from_user("Hi"), + ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="mytool", arguments={"a": 1}, id="123")]), + ] + history = ChatHistory(messages=messages) + assert history.model_dump() == {"messages": [message.to_dict() for message in messages]} + assert history.model_dump(mode="json") == {"messages": [message.to_dict() for message in messages]} + + def test_model_validate_serialized_messages(self): + messages = [ChatMessage.from_user("Hi"), ChatMessage.from_assistant("Hello")] + history = ChatHistory.model_validate({"messages": [message.to_dict() for message in messages]}) + assert history.messages == messages + + def test_instances_pass_through_unchanged(self): + message = ChatMessage.from_user("Hi") + history = ChatHistory(messages=[message]) + assert history.messages[0] is message + + def test_round_trip(self, base64_image_string): + tool_call = ToolCall(id="123", tool_name="mytool", arguments={"a": 1}) + messages = [ + ChatMessage.from_system("You are helpful."), + ChatMessage( + _role=ChatRole.ASSISTANT, + _content=[ReasoningContent(reasoning_text="Thinking..."), TextContent(text="Checking."), tool_call], + _meta={"some": "info"}, + ), + ChatMessage.from_tool(tool_result="42", origin=tool_call), + ChatMessage( + _role=ChatRole.USER, + _content=[ + ImageContent(base64_image=base64_image_string, mime_type="image/png"), + FileContent(base64_data="aGVsbG8=", mime_type="text/plain", filename="hello.txt"), + ], + ), + ] + + dumped = ChatHistory(messages=messages).model_dump(mode="json") + assert ChatHistory.model_validate(dumped).messages == messages + + def test_invalid_value_raises_validation_error(self): + with pytest.raises(ValidationError): + ChatHistory(messages=["not a message"]) # type: ignore[list-item] + + class TestToOpenaiDictFormat: def test_to_openai_dict_format_system_message(self): message = ChatMessage.from_system("You are good assistant")