codekingpro/portable-devtools

1"""Chat Message."""2 3from typing import Any, Literal4 5from typing_extensions import override6 7from langchain_core.messages.base import (8    BaseMessage,9    BaseMessageChunk,10    merge_content,11)12from langchain_core.utils._merge import merge_dicts13 14 15class ChatMessage(BaseMessage):16    """Message that can be assigned an arbitrary speaker (i.e. role)."""17 18    role: str19    """The speaker / role of the Message."""20 21    type: Literal["chat"] = "chat"22    """The type of the message (used during serialization)."""23 24 25class ChatMessageChunk(ChatMessage, BaseMessageChunk):26    """Chat Message chunk."""27 28    # Ignoring mypy re-assignment here since we're overriding the value29    # to make sure that the chunk variant can be discriminated from the30    # non-chunk variant.31    type: Literal["ChatMessageChunk"] = "ChatMessageChunk"  # type: ignore[assignment]32    """The type of the message (used during serialization)."""33 34    @override35    def __add__(self, other: Any) -> BaseMessageChunk:  # type: ignore[override]36        if isinstance(other, ChatMessageChunk):37            if self.role != other.role:38                msg = "Cannot concatenate ChatMessageChunks with different roles."39                raise ValueError(msg)40 41            return self.__class__(42                role=self.role,43                content=merge_content(self.content, other.content),44                additional_kwargs=merge_dicts(45                    self.additional_kwargs, other.additional_kwargs46                ),47                response_metadata=merge_dicts(48                    self.response_metadata, other.response_metadata49                ),50                id=self.id,51            )52        if isinstance(other, BaseMessageChunk):53            return self.__class__(54                role=self.role,55                content=merge_content(self.content, other.content),56                additional_kwargs=merge_dicts(57                    self.additional_kwargs, other.additional_kwargs58                ),59                response_metadata=merge_dicts(60                    self.response_metadata, other.response_metadata61                ),62                id=self.id,63            )64        return super().__add__(other)65