Moved constructor args from Archivist to Speaker when it made sense

Added a periodic save time for chats in memory
Added flags for mute time, save time, chats folder
This commit is contained in:
vylion 2020-10-27 17:10:41 +01:00
parent 895e1cd843
commit 0974ec52e7
4 changed files with 94 additions and 55 deletions

View file

@ -7,8 +7,8 @@ from generator import Generator
class Archivist(object): class Archivist(object):
def __init__(self, logger, chatdir=None, chatext=None, admin=0, def __init__(self, logger, chatdir=None, chatext=None, admin=0,
period_inc=5, save_count=15, max_period=100000, max_len=50, period_inc=5, save_count=15, max_period=100000,
read_only=False, filter_cids=None, bypass=False read_only=False
): ):
if chatdir is None or len(chatdir) == 0: if chatdir is None or len(chatdir) == 0:
chatdir = "./" chatdir = "./"
@ -17,20 +17,16 @@ class Archivist(object):
self.logger = logger self.logger = logger
self.chatdir = chatdir self.chatdir = chatdir
self.chatext = chatext self.chatext = chatext
self.admin = admin
self.period_inc = period_inc self.period_inc = period_inc
self.save_count = save_count self.save_count = save_count
self.max_period = max_period self.max_period = max_period
self.max_len = max_len
self.read_only = read_only self.read_only = read_only
self.filter_cids = filter_cids
self.bypass = bypass
def chat_folder(self, *formatting, **key_format): 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): 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): def store(self, tag, data, gen):
chat_folder = self.chat_folder(tag=tag) chat_folder = self.chat_folder(tag=tag)
@ -121,10 +117,7 @@ class Archivist(object):
reader = self.get_reader(cid) reader = self.get_reader(cid)
# self.logger.info("Chat {} contents:\n{}".format(cid, reader.card.dumps())) # self.logger.info("Chat {} contents:\n{}".format(cid, reader.card.dumps()))
self.logger.info("Successfully passed through {} ({}) chat.\n".format(cid, reader.title())) self.logger.info("Successfully passed through {} ({}) chat.\n".format(cid, reader.title()))
if self.bypass: # I forgot what I made this for if reader.period() > self.max_period:
reader.set_period(random.randint(self.max_period // 2, self.max_period))
self.store(*reader.archive())
elif reader.period() > self.max_period:
reader.set_period(self.max_period) reader.set_period(self.max_period)
self.store(*reader.archive()) self.store(*reader.archive())
yield reader yield reader

View file

@ -64,3 +64,6 @@ class MemoryList(MutableSequence):
self._list.remove(val) self._list.remove(val)
self._list.append(val) self._list.append(val)
return val return val
def remove(self, val):
self._list.remove(val)

View file

@ -2,13 +2,14 @@
import random import random
import time import time
from sys import stderr
from memorylist import MemoryList from memorylist import MemoryList
from reader import Reader, get_chat_title from reader import Reader, get_chat_title
from telegram.error import TimedOut, NetworkError from telegram.error import NetworkError
def eprint(*args, **kwargs): 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): 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): if text.startswith(Reader.TAG_PREFIX):
words = text.split(maxsplit=1) words = text.split(maxsplit=1)
if logger: if logger:
# logger.info('Sending {} "{}" to {}'.format(words[0][4:-1], words[1], cid)) logger.info('Sending {} "{}" to {}'.format(words[0][4:-1], words[1], cid))
eprint('.') # eprint('[]')
# Logs something like 'Sending VIDEO "VIDEO_ID" to CHAT_ID' # Logs something like 'Sending VIDEO "VIDEO_ID" to CHAT_ID'
if words[0] == Reader.STICKER_TAG: if words[0] == Reader.STICKER_TAG:
@ -33,37 +34,48 @@ def send(bot, cid, text, replying=None, formatting=None, logger=None, **kwargs):
if logger: if logger:
mtype = "reply" if replying else "message" mtype = "reply" if replying else "message"
logger.info("Sending a {} to {}: '{}'".format(mtype, cid, text)) logger.info("Sending a {} to {}: '{}'".format(mtype, cid, text))
# eprint('.')
return bot.send_message(cid, text, **kwargs) return bot.send_message(cid, text, **kwargs)
class Speaker(object): class Speaker(object):
ModeFixed = "FIXED_MODE" 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, 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.names = nicknames
self.mute_time = mute_time self.mute_time = mute_time
self.mute_timer = None
self.username = username 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("----")
logger.info("Finished loading.") logger.info("Finished loading.")
logger.info("Loaded {} chats.".format(archivist.chat_count())) logger.info("Loaded {} chats.".format(archivist.chat_count()))
logger.info("----") logger.info("----")
self.wakeup = wakeup self.wakeup = wakeup
self.logger = logger self.logger = logger
self.reply = reply self.reply = reply
self.repeat = repeat self.repeat = repeat
self.filter_cids = archivist.filter_cids self.filter_cids = filter_cids
self.bypass = archivist.bypass
self.time_counter = None
self.memory = MemoryList(memory) 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)): def announce(self, bot, announcement, check=(lambda _: True)):
# Sends an announcement to all chats that pass the check # Sends an announcement to all chats that pass the check
for reader in self.archivist.readers_pass(): for reader in self.readers_pass():
try: try:
if check(reader): if check(reader):
send(bot, reader.cid(), announcement) send(bot, reader.cid(), announcement)
@ -74,7 +86,7 @@ class Speaker(object):
def wake(self, bot, wake): def wake(self, bot, wake):
# If wakeup flag is set, sends a wake-up message as announcement to all chats that # 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' # 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: if self.wakeup:
def group_check(reader): def group_check(reader):
@ -90,9 +102,9 @@ class Speaker(object):
if reader is not None: if reader is not None:
return reader return reader
reader = self.archivist.get_reader(cid) reader = self.get_reader_file(cid)
if not reader: 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) old_reader = self.memory.append(reader)
if old_reader is not None: if old_reader is not None:
@ -104,7 +116,7 @@ class Speaker(object):
def access_reader(self, cid): def access_reader(self, cid):
reader = self.get_reader(cid) reader = self.get_reader(cid)
if reader is None: if reader is None:
return self.archivist.get_reader(cid) return self.get_reader_file(cid)
return reader return reader
def mentioned(self, text): def mentioned(self, text):
@ -117,7 +129,7 @@ class Speaker(object):
def is_mute(self): def is_mute(self):
current_time = int(time.perf_counter()) 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): def should_reply(self, message, reader):
if self.is_mute(): if self.is_mute():
@ -136,9 +148,25 @@ class Speaker(object):
if reader is None: if reader is None:
raise ValueError("Tried to store a None Reader.") raise ValueError("Tried to store a None Reader.")
else: 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): def read(self, update, context):
self.save()
if update.message is None: if update.message is None:
return return
chat = update.message.chat chat = update.message.chat
@ -178,18 +206,19 @@ class Speaker(object):
self.say(context.bot, reader, replying=rid) self.say(context.bot, reader, replying=rid)
def user_is_admin(self, member): 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("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))) # eprint('!')
# self.logger.info("Bot Creator ID is {}".format(str(self.admin)))
return ((member.status == 'creator') return ((member.status == 'creator')
or (member.status == 'administrator') or (member.status == 'administrator')
or (member.user.id == self.archivist.admin)) or (member.user.id == self.admin))
def speech(self, reader): 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): def say(self, bot, reader, replying=None, **kwargs):
cid = reader.cid() 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 return
if self.is_mute(): if self.is_mute():
return return
@ -197,17 +226,14 @@ class Speaker(object):
try: try:
send(bot, cid, self.speech(reader), replying, logger=self.logger, **kwargs) send(bot, cid, self.speech(reader), replying, logger=self.logger, **kwargs)
if self.bypass: if self.bypass:
max_period = self.archivist.max_period max_period = self.max_period
reader.set_period(random.randint(max_period // 4, max_period)) reader.set_period(random.randint(max_period // 4, max_period))
if random.random() <= self.repeat: if random.random() <= self.repeat:
send(bot, cid, self.speech(reader), logger=self.logger, **kwargs) 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: except NetworkError as e:
if '429' in e.message: if '429' in e.message:
self.logger.error("Error: TooManyRequests. Going mute for {} seconds.".format(self.mute_time)) 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: else:
self.logger.error("Sending a message caused network error:") self.logger.error("Sending a message caused network error:")
self.logger.exception(e) self.logger.exception(e)
@ -223,7 +249,7 @@ class Speaker(object):
update.message.reply_text("I remember {} messages.".format(num)) update.message.reply_text("I remember {} messages.".format(num))
def get_chats(self, update, context): 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) chat_list = "\n".join(lines)
update.message.reply_text("I have the following chats:\n\n" + chat_list) 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 = int(words[1])
period = reader.set_period(period) period = reader.set_period(period)
update.message.reply_text("Period of speaking set to {}.".format(period)) update.message.reply_text("Period of speaking set to {}.".format(period))
self.archivist.store(*reader.archive()) self.store_file(*reader.archive())
except Exception: except Exception:
update.message.reply_text("Format was confusing; period unchanged from {}.".format(reader.period())) 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 = float(words[1])
answer = reader.set_answer(answer) answer = reader.set_answer(answer)
update.message.reply_text("Answer probability set to {}.".format(answer)) update.message.reply_text("Answer probability set to {}.".format(answer))
self.archivist.store(*reader.archive()) self.store_file(*reader.archive())
except Exception: except Exception:
update.message.reply_text("Format was confusing; answer probability unchanged from {}.".format(reader.answer())) update.message.reply_text("Format was confusing; answer probability unchanged from {}.".format(reader.answer()))
@ -286,7 +312,7 @@ class Speaker(object):
reader.toggle_restrict() reader.toggle_restrict()
allowed = "let only admins" if reader.is_restricted() else "let everyone" allowed = "let only admins" if reader.is_restricted() else "let everyone"
update.message.reply_text("I will {} configure me now.".format(allowed)) 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): def silence(self, update, context):
if "group" not in update.message.chat.type: if "group" not in update.message.chat.type:
@ -303,7 +329,7 @@ class Speaker(object):
reader.toggle_silence() reader.toggle_silence()
allowed = "avoid mentioning" if reader.is_silenced() else "mention" allowed = "avoid mentioning" if reader.is_silenced() else "mention"
update.message.reply_text("I will {} people now.".format(allowed)) 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): def who(self, update, context):
msg = update.message msg = update.message

View file

@ -73,11 +73,22 @@ def stop(update, context):
def main(): def main():
global speakerbot global speakerbot
parser = argparse.ArgumentParser(description='A Telegram markov bot.') 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('token', metavar='TOKEN',
parser.add_argument('admin_id', metavar='ADMIN_ID', type=int, help='The ID of the Telegram user that manages this bot') help='The Bot Token to work with the Telegram Bot API')
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('admin_id', metavar='ADMIN_ID', type=int, default=0,
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)') help='The ID of the Telegram user that manages this bot')
parser.add_argument('-n', '--nicknames', nargs='*', metavar='NICKNAME', help='Any possible nicknames that the bot could answer to.') 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() args = parser.parse_args()
@ -89,15 +100,21 @@ def main():
filter_cids.append(str(args.admin_id)) filter_cids.append(str(args.admin_id))
archivist = Archivist(logger, archivist = Archivist(logger,
chatdir="./chatlogs/", chatdir=args.directory,
chatext=".vls", chatext=".vls",
admin=args.admin_id,
filter_cids=filter_cids,
read_only=False read_only=False
) )
username = updater.bot.get_me().username 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 # Get the dispatcher to register handlers
dp = updater.dispatcher dp = updater.dispatcher
@ -109,7 +126,7 @@ def main():
dp.add_handler(CommandHandler("help", static_reply(help_msg))) dp.add_handler(CommandHandler("help", static_reply(help_msg)))
dp.add_handler(CommandHandler("count", speakerbot.get_count)) dp.add_handler(CommandHandler("count", speakerbot.get_count))
dp.add_handler(CommandHandler("period", speakerbot.period)) 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("user", get_name, Filters.chat(chat_id=archivist.admin)))
# dp.add_handler(CommandHandler("id", get_id)) # dp.add_handler(CommandHandler("id", get_id))
dp.add_handler(CommandHandler("stop", stop)) dp.add_handler(CommandHandler("stop", stop))