diff --git a/archivist.py b/archivist.py index 80cd873..29d72db 100644 --- a/archivist.py +++ b/archivist.py @@ -7,8 +7,8 @@ from generator import Generator class Archivist(object): def __init__(self, logger, chatdir=None, chatext=None, admin=0, - period_inc=5, save_count=15, max_period=100000, max_len=50, - read_only=False, filter_cids=None, bypass=False + period_inc=5, save_count=15, max_period=100000, + read_only=False ): if chatdir is None or len(chatdir) == 0: chatdir = "./" @@ -17,20 +17,16 @@ class Archivist(object): self.logger = logger self.chatdir = chatdir self.chatext = chatext - self.admin = admin self.period_inc = period_inc self.save_count = save_count self.max_period = max_period - self.max_len = max_len self.read_only = read_only - self.filter_cids = filter_cids - self.bypass = bypass def chat_folder(self, *formatting, **key_format): - return (self.chatdir + "chat_{tag}").format(*formatting, **key_format) + return (self.chatdir + "/chat_{tag}").format(*formatting, **key_format) def chat_file(self, *formatting, **key_format): - return (self.chatdir + "chat_{tag}/{file}{ext}").format(*formatting, **key_format) + return (self.chatdir + "/chat_{tag}/{file}{ext}").format(*formatting, **key_format) def store(self, tag, data, gen): chat_folder = self.chat_folder(tag=tag) @@ -121,10 +117,7 @@ class Archivist(object): reader = self.get_reader(cid) # self.logger.info("Chat {} contents:\n{}".format(cid, reader.card.dumps())) self.logger.info("Successfully passed through {} ({}) chat.\n".format(cid, reader.title())) - if self.bypass: # I forgot what I made this for - reader.set_period(random.randint(self.max_period // 2, self.max_period)) - self.store(*reader.archive()) - elif reader.period() > self.max_period: + if reader.period() > self.max_period: reader.set_period(self.max_period) self.store(*reader.archive()) yield reader diff --git a/memorylist.py b/memorylist.py index cd1afa4..f62c05f 100644 --- a/memorylist.py +++ b/memorylist.py @@ -64,3 +64,6 @@ class MemoryList(MutableSequence): self._list.remove(val) self._list.append(val) return val + + def remove(self, val): + self._list.remove(val) diff --git a/speaker.py b/speaker.py index 5810e7c..dc3b20f 100644 --- a/speaker.py +++ b/speaker.py @@ -2,13 +2,14 @@ import random import time +from sys import stderr from memorylist import MemoryList from reader import Reader, get_chat_title -from telegram.error import TimedOut, NetworkError +from telegram.error import NetworkError def eprint(*args, **kwargs): - print(*args, file=sys.stderr, **kwargs) + print(*args, end=' ', file=stderr, **kwargs) def send(bot, cid, text, replying=None, formatting=None, logger=None, **kwargs): @@ -18,8 +19,8 @@ def send(bot, cid, text, replying=None, formatting=None, logger=None, **kwargs): if text.startswith(Reader.TAG_PREFIX): words = text.split(maxsplit=1) if logger: - # logger.info('Sending {} "{}" to {}'.format(words[0][4:-1], words[1], cid)) - eprint('.') + logger.info('Sending {} "{}" to {}'.format(words[0][4:-1], words[1], cid)) + # eprint('[]') # Logs something like 'Sending VIDEO "VIDEO_ID" to CHAT_ID' if words[0] == Reader.STICKER_TAG: @@ -33,37 +34,48 @@ def send(bot, cid, text, replying=None, formatting=None, logger=None, **kwargs): if logger: mtype = "reply" if replying else "message" logger.info("Sending a {} to {}: '{}'".format(mtype, cid, text)) + # eprint('.') return bot.send_message(cid, text, **kwargs) class Speaker(object): ModeFixed = "FIXED_MODE" - ModeChance = "MODE_CHANCE" + ModeChance = "CHANCE_MODE" - def __init__(self, username, archivist, logger, nicknames=[], mute_time=60, + def __init__(self, username, archivist, logger, admin=0, nicknames=[], reply=0.1, repeat=0.05, wakeup=False, mode=ModeFixed, - memory=20 + memory=20, mute_time=60, save_time=3600, bypass=False, + filter_cids=[], max_len=50 ): self.names = nicknames self.mute_time = mute_time + self.mute_timer = None self.username = username - self.archivist = archivist + + self.max_period = archivist.max_period + self.get_reader_file = archivist.get_reader + self.store_file = archivist.store + self.readers_pass = archivist.readers_pass + logger.info("----") logger.info("Finished loading.") logger.info("Loaded {} chats.".format(archivist.chat_count())) logger.info("----") + self.wakeup = wakeup self.logger = logger self.reply = reply self.repeat = repeat - self.filter_cids = archivist.filter_cids - self.bypass = archivist.bypass - self.time_counter = None + self.filter_cids = filter_cids self.memory = MemoryList(memory) + self.memory_timer = time.perf_counter() + self.admin = admin + self.bypass = bypass + self.max_len = max_len def announce(self, bot, announcement, check=(lambda _: True)): # Sends an announcement to all chats that pass the check - for reader in self.archivist.readers_pass(): + for reader in self.readers_pass(): try: if check(reader): send(bot, reader.cid(), announcement) @@ -74,7 +86,7 @@ class Speaker(object): def wake(self, bot, wake): # If wakeup flag is set, sends a wake-up message as announcement to all chats that # are groups. Also, always sends a wakeup message to the 'bot admin' - send(bot, self.archivist.admin, wake) + send(bot, self.admin, wake) if self.wakeup: def group_check(reader): @@ -90,9 +102,9 @@ class Speaker(object): if reader is not None: return reader - reader = self.archivist.get_reader(cid) + reader = self.get_reader_file(cid) if not reader: - reader = Reader.FromChat(chat, self.archivist.max_period, self.logger) + reader = Reader.FromChat(chat, self.max_period, self.logger) old_reader = self.memory.append(reader) if old_reader is not None: @@ -104,7 +116,7 @@ class Speaker(object): def access_reader(self, cid): reader = self.get_reader(cid) if reader is None: - return self.archivist.get_reader(cid) + return self.get_reader_file(cid) return reader def mentioned(self, text): @@ -117,7 +129,7 @@ class Speaker(object): def is_mute(self): current_time = int(time.perf_counter()) - return self.time_counter is not None and (current_time - self.time_counter) < self.mute_time + return self.mute_timer is not None and (current_time - self.mute_timer) < self.mute_time def should_reply(self, message, reader): if self.is_mute(): @@ -136,9 +148,25 @@ class Speaker(object): if reader is None: raise ValueError("Tried to store a None Reader.") else: - self.archivist.store(*reader.archive()) + self.store_file(*reader.archive()) + + def should_save(self): + current_time = int(time.perf_counter()) + elapsed = (current_time - self.memory_timer) + self.logger.debug("Save check: {}".format(elapsed)) + return elapsed < self.save_time + + def save(self): + if self.should_save(): + self.logger.info("Saving chats in memory...") + for reader in self.memory: + self.store(reader) + self.memory_timer = time.perf_counter() + self.logger.info("Chats saved.") def read(self, update, context): + self.save() + if update.message is None: return chat = update.message.chat @@ -178,18 +206,19 @@ class Speaker(object): self.say(context.bot, reader, replying=rid) def user_is_admin(self, member): - # self.logger.info("user {} ({}) requesting a restricted action".format(str(member.user.id), member.user.name)) - # self.logger.info("Bot Creator ID is {}".format(str(self.archivist.admin))) + self.logger.info("user {} ({}) requesting a restricted action".format(str(member.user.id), member.user.name)) + # eprint('!') + # self.logger.info("Bot Creator ID is {}".format(str(self.admin))) return ((member.status == 'creator') or (member.status == 'administrator') - or (member.user.id == self.archivist.admin)) + or (member.user.id == self.admin)) def speech(self, reader): - return reader.generate_message(self.archivist.max_len) + return reader.generate_message(self.max_len) def say(self, bot, reader, replying=None, **kwargs): cid = reader.cid() - if self.filter_cids is not None and cid not in self.filter_cids: + if cid not in self.filter_cids: return if self.is_mute(): return @@ -197,17 +226,14 @@ class Speaker(object): try: send(bot, cid, self.speech(reader), replying, logger=self.logger, **kwargs) if self.bypass: - max_period = self.archivist.max_period + max_period = self.max_period reader.set_period(random.randint(max_period // 4, max_period)) if random.random() <= self.repeat: send(bot, cid, self.speech(reader), logger=self.logger, **kwargs) - except TimedOut as e: - self.logger.error("Telegram timed out.") - self.logger.exception(e) except NetworkError as e: if '429' in e.message: self.logger.error("Error: TooManyRequests. Going mute for {} seconds.".format(self.mute_time)) - self.time_counter = int(time.perf_counter()) + self.mute_timer = int(time.perf_counter()) else: self.logger.error("Sending a message caused network error:") self.logger.exception(e) @@ -223,7 +249,7 @@ class Speaker(object): update.message.reply_text("I remember {} messages.".format(num)) def get_chats(self, update, context): - lines = ["[{}]: {}".format(reader.cid(), reader.title()) for reader in self.archivist.readers_pass] + lines = ["[{}]: {}".format(reader.cid(), reader.title()) for reader in self.readers_pass()] chat_list = "\n".join(lines) update.message.reply_text("I have the following chats:\n\n" + chat_list) @@ -245,7 +271,7 @@ class Speaker(object): period = int(words[1]) period = reader.set_period(period) update.message.reply_text("Period of speaking set to {}.".format(period)) - self.archivist.store(*reader.archive()) + self.store_file(*reader.archive()) except Exception: update.message.reply_text("Format was confusing; period unchanged from {}.".format(reader.period())) @@ -267,7 +293,7 @@ class Speaker(object): answer = float(words[1]) answer = reader.set_answer(answer) update.message.reply_text("Answer probability set to {}.".format(answer)) - self.archivist.store(*reader.archive()) + self.store_file(*reader.archive()) except Exception: update.message.reply_text("Format was confusing; answer probability unchanged from {}.".format(reader.answer())) @@ -286,7 +312,7 @@ class Speaker(object): reader.toggle_restrict() allowed = "let only admins" if reader.is_restricted() else "let everyone" update.message.reply_text("I will {} configure me now.".format(allowed)) - self.archivist.store(*reader.archive()) + self.store_file(*reader.archive()) def silence(self, update, context): if "group" not in update.message.chat.type: @@ -303,7 +329,7 @@ class Speaker(object): reader.toggle_silence() allowed = "avoid mentioning" if reader.is_silenced() else "mention" update.message.reply_text("I will {} people now.".format(allowed)) - self.archivist.store(*reader.archive()) + self.store_file(*reader.archive()) def who(self, update, context): msg = update.message diff --git a/velasco.py b/velasco.py index b9d1953..48341d2 100644 --- a/velasco.py +++ b/velasco.py @@ -73,11 +73,22 @@ def stop(update, context): def main(): global speakerbot parser = argparse.ArgumentParser(description='A Telegram markov bot.') - parser.add_argument('token', metavar='TOKEN', help='The Bot Token to work with the Telegram Bot API') - parser.add_argument('admin_id', metavar='ADMIN_ID', type=int, help='The ID of the Telegram user that manages this bot') - parser.add_argument('-w', '--wakeup', action='store_true', help='Flag that makes the bot send a first message to all chats during wake up.') - parser.add_argument('-f', '--filter', nargs='*', metavar='FILTER_CID', help='Zero or more chat IDs to add in a filter whitelist (default is empty, all chats allowed)') - parser.add_argument('-n', '--nicknames', nargs='*', metavar='NICKNAME', help='Any possible nicknames that the bot could answer to.') + parser.add_argument('token', metavar='TOKEN', + help='The Bot Token to work with the Telegram Bot API') + parser.add_argument('admin_id', metavar='ADMIN_ID', type=int, default=0, + help='The ID of the Telegram user that manages this bot') + parser.add_argument('-w', '--wakeup', action='store_true', + help='Flag that makes the bot send a first message to all chats during wake up.') + parser.add_argument('-f', '--filter', nargs='*', default=[], metavar='cid', + help='Zero or more chat IDs to add in a filter whitelist (default is empty, all chats allowed)') + parser.add_argument('-n', '--nicknames', nargs='*', default=[], metavar='name', + help='Any possible nicknames that the bot could answer to.') + parser.add_argument('-d', '--directory', metavar='CHATLOG_DIR', default='./chatlogs', + help='The chat logs directory path (default: "./chatlogs").') + parser.add_argument('-m', '--mute_time', metavar='T', type=int, default=60, + help='The time (in s) for the muting period when Telegram limits the bot. (default: 60).') + parser.add_argument('-s', '--save_time', metavar='T', type=int, default=3600, + help='The time (in s) for periodic saves (default: 3600).') args = parser.parse_args() @@ -89,15 +100,21 @@ def main(): filter_cids.append(str(args.admin_id)) archivist = Archivist(logger, - chatdir="./chatlogs/", + chatdir=args.directory, chatext=".vls", - admin=args.admin_id, - filter_cids=filter_cids, read_only=False ) username = updater.bot.get_me().username - speakerbot = Speaker("@" + username, archivist, logger, nicknames=args.nicknames, wakeup=args.wakeup) + speakerbot = Speaker("@" + username, + archivist, + logger, + admin=args.admin_id, + filter_cids=filter_cids, + nicknames=args.nicknames, + wakeup=args.wakeup, + mute_time=args.mute_time, + save_time=args.save_time) # Get the dispatcher to register handlers dp = updater.dispatcher @@ -109,7 +126,7 @@ def main(): dp.add_handler(CommandHandler("help", static_reply(help_msg))) dp.add_handler(CommandHandler("count", speakerbot.get_count)) dp.add_handler(CommandHandler("period", speakerbot.period)) - dp.add_handler(CommandHandler("list", speakerbot.get_chats, filters=Filters.chat(chat_id=archivist.admin))) + dp.add_handler(CommandHandler("list", speakerbot.get_chats, filters=Filters.chat(chat_id=speakerbot.admin))) # dp.add_handler(CommandHandler("user", get_name, Filters.chat(chat_id=archivist.admin))) # dp.add_handler(CommandHandler("id", get_id)) dp.add_handler(CommandHandler("stop", stop))