matrix-pygmalion-bot/matrix_pygmalion_bot/chatlog.py

import os
import matrix_pygmalion_bot.translate as translate
import logging

logger = logging.getLogger(__name__)


class ChatMessage:
    def __init__(self, event_id, timestamp, user_name, is_own_message, is_command, relates_to_event, message, language="en", english_original_message=None):
        self.event_id = event_id
        self.timestamp = timestamp
        self.user_name = user_name
        self.is_own_message = is_own_message
        self.is_command = is_command
        self.relates_to_event = relates_to_event
        self.num_tokens = None
        self.message = {}
        self.message[language] = message
        if not (language == "en"):
            if not (english_original_message is None):
                self.message["en"] = english_original_message
            else:
                self.message["en"] = translate.translate(message, language, "en")
        self.language = language
        self.num_tokens = None

    def __str__(self):
        return str("{}: {}".format(self.user_name, self.message[self.language]))
    def getTranslation(self, to_lang):
        if not (to_lang in self.message):
            self.message[to_lang] = translate.translate(self.message["en"], "en", to_lang)
        return self.message[to_lang]


class ChatHistory:
    def __init__(self, bot_name, room_name):
        self.bot_name = bot_name
        self.room_name = room_name
        self.chat_history = {}
    def __str__(self):
        return str("Chat History for {} in room {}".format(self.bot_name, self.room_name))
    def getLen(self):
        return len(self.chat_history)
    def load_from_file(self):
        pass
    def clear(self):
        self.chat_history = {}
    def remove(self, num_items=1):
        for i in range(num_items):
            event_id, item = self.chat_history.popitem()
        return item        
    def add(self, event_id, timestamp, user_name, is_own_message, is_command, relates_to_event, message, language="en", english_original_message=None):
        chat_message = ChatMessage(event_id, timestamp, user_name, is_own_message, is_command, relates_to_event, message, language, english_original_message)
        self.chat_history[chat_message.event_id] = chat_message
        os.makedirs("./chatlogs", exist_ok=True)
        with open("chatlogs/" + self.bot_name + "_" + self.room_name + ".txt", "a") as f:
            f.write("{}: {}\n".format(user_name, message))
        return chat_message
    def getLastItem(self):
        key, chat_item = list(self.chat_history.items())[-1]
        return chat_item

class BotChatHistory:
    def __init__(self, bot_name):
        self.bot_name = bot_name
        self.chat_rooms = {}
    def __str__(self):
        return str("Chat History for {}".format(self.bot_name))
    def room (self, room):
        if not room in self.chat_rooms:
            self.chat_rooms[room] = ChatHistory(self.bot_name, room)
        return self.chat_rooms[room]