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, role=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
        self.role = role

    @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):
        if self.role == "human":
            return HumanMessage(
                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,
                }
            )
        elif self.role == "ai":
            return AIMessage(
                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,
                }
            )
        elif self.role == "system":
            return SystemMessage(
                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,
                }
            )
        else:
            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):
        if self.message.startswith('<ERROR>'):
            return True
        elif self.message.startswith('<WARNING>'):
            return True
        else:
            return False

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