diff --git a/matrix_pygmalion_bot/ai/runpod_pygmalion.py b/matrix_pygmalion_bot/ai/runpod_pygmalion.py index 550bb4b..4791c78 100644 --- a/matrix_pygmalion_bot/ai/runpod_pygmalion.py +++ b/matrix_pygmalion_bot/ai/runpod_pygmalion.py @@ -82,7 +82,7 @@ async def generate_sync( reply = answer[:idx].strip() else: reply = answer.removesuffix('<|endoftext|>').strip() - reply.replace("\n{bot_name}: ", " ") + reply = reply.replace("\n{bot_name}: ", " ") return reply else: return "" diff --git a/matrix_pygmalion_bot/chatlog.py b/matrix_pygmalion_bot/chatlog.py index a9fc59a..f827822 100644 --- a/matrix_pygmalion_bot/chatlog.py +++ b/matrix_pygmalion_bot/chatlog.py @@ -2,11 +2,12 @@ import os import matrix_pygmalion_bot.translate as translate class ChatMessage: - def __init__(self, event_id, timestamp, user_name, is_own_message, relates_to_event, message, language="en", english_original_message=None): + 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 = {} @@ -18,6 +19,7 @@ class ChatMessage: 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): @@ -43,15 +45,15 @@ class ChatHistory: 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, relates_to_event, message, language="en", english_original_message=None): - chat_message = ChatMessage(event_id, timestamp, user_name, is_own_message, relates_to_event, message, language, english_original_message) + 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 = reversed(self.chat_history.items())[0] + key, chat_item = list(self.chat_history.items())[-1] return chat_item class BotChatHistory: diff --git a/matrix_pygmalion_bot/core.py b/matrix_pygmalion_bot/core.py index cc46482..6e45b66 100644 --- a/matrix_pygmalion_bot/core.py +++ b/matrix_pygmalion_bot/core.py @@ -38,28 +38,22 @@ class Callbacks(object): async def message_cb(self, room: MatrixRoom, event: RoomMessageText) -> None: if not hasattr(event, 'body'): return - #message = event.body - is_own_message = False - if event.sender == self.client.user: - is_own_message = True - is_command = False - if event.body.startswith('!'): - is_command = True - if not (self.bot.owner is None): - if not (event.sender == self.bot.owner or is_own_message): - return relates_to = None if 'm.relates_to' in event.source["content"]: relates_to = event.source["content"]['m.relates_to']["event_id"] + is_command = False + if event.body.startswith('!'): + is_command = True language = "en" if not (self.bot.translate is None) and not is_command: language = self.bot.translate + if 'original_message' in event.source["content"]: english_original_message = event.source["content"]['original_message'] else: english_original_message = None - chat_message = self.bot.chat_history.room(room.display_name).add(event.event_id, event.server_timestamp, room.user_name(event.sender), is_own_message, relates_to, event.body, language, english_original_message) + chat_message = self.bot.chat_history.room(room.display_name).add(event.event_id, event.server_timestamp, room.user_name(event.sender), event.sender == self.client.user, is_command, relates_to, event.body, language, english_original_message) if self.bot.not_synced: return print( @@ -70,11 +64,12 @@ class Callbacks(object): await self.client.room_read_markers(room.room_id, event.event_id, event.event_id) # Ignore messages from ourselves - if is_own_message: - return - - if not hasattr(event, 'body'): + if chat_message.is_own_message: return + # Ignore message from strangers + if not (self.bot.owner is None): + if not (event.sender == self.bot.owner or chat_message.is_own_message): + return if event.body.startswith('!replybot'): print(event)