matrix-pygmalion-bot/matrix_pygmalion_bot/bot/utilities/messages.py

from langchain.schema import AIMessage, HumanMessage, SystemMessage, ChatMessage


class Message(object):
    def __init__(self, timestamp, user_name, message, event_id=None, user_id=None, room_name=None, room_id=None):
        self.timestamp = timestamp
        self.user_name = user_name
        self.message = message
        self.event_id = event_id
        self.user_id = user_id
        self.room_name = room_name
        self.room_id = room_id

    @classmethod
    def from_matrix(cls, room, event):
        return cls(event.server_timestamp/1000, room.user_name(event.sender), event.body, event.event_id, event.sender, room.display_name, room.room_id)

    def to_langchain(self):
        return ChatMessage(
            content=self.message,
            role=self.user_name, # "chat"
            additional_kwargs={
                "timestamp": self.timestamp,
                "user_name": self.user_name,
                "event_id": self.event_id,
                "user_id": self.user_id,
                "room_name": self.room_name,
                "room_id": self.room_id,
            }
        )

    def is_from(self, user_id):
        return self.user_id == user_id

    def is_command(self):
        return self.message.startswith('!')

    def is_error(self):
        return self.message.startswith('<ERROR>')

    def __str__(self):
        return str("{}: {}".format(self.user_name, self.message))