You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
97 lines
3.9 KiB
97 lines
3.9 KiB
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]
|
|
|