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.stop_here = False
        self.relates_to_event = relates_to_event
        self.num_tokens = None
        self.is_in_saved_prompt = False
        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]
    def updateText(self, new_text, language="en"):
        self.message[self.language] = new_text
        self.num_tokens = None
        if not (language == "en"):
            self.message["en"] = translate.translate(message, language, "en")

class ChatHistory:
    def __init__(self, bot_name, room_name):
        self.bot_name = bot_name
        self.room_name = room_name
        self.context_fast_forward = False
        self.saved_context = None
        self.saved_context_num_tokens = None
        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 remove_id(self, event_id):
        if event_id in self.chat_history:
            del self.chat_history[event_id]
        else:
            logger.warning("remove_id: could not delete event. event_id not found in chat history")
    def exists_id(self, event_id):
        return (event_id in self.chat_history)
    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
    def setFastForward(self, value):
        self.context_fast_forward = value
    def getFastForward(self):
        return self.context_fast_forward
    def getSavedPrompt(self):
        return self.saved_context
    def setSavedPrompt(self, context, num_tokens):
        self.saved_context = context
        self.saved_context_num_tokens = num_tokens

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]