diff --git a/.gitignore b/.gitignore index 4ccade3..1ae7c90 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ chatlogs/* __pycache__/* misc/* +bkp/* test/* diff --git a/archivist.py b/archivist.py index 21637a1..9ff340e 100644 --- a/archivist.py +++ b/archivist.py @@ -1,61 +1,97 @@ -import os, errno, random, pickle -from chatreader import ChatReader as Reader +import os, random, pickle +from reader import Reader from generator import Generator class Archivist(object): def __init__(self, logger, chatdir=None, chatext=None, admin=0, - freq_increment=5, save_count=15, max_period=100000, max_len=50, + period_inc=5, save_count=15, max_period=100000, max_len=50, read_only=False, filter_cids=None, bypass=False ): if chatdir is None or len(chatdir) == 0: - raise ValueError("Chatlog directory name is empty") - elif chatext is None: # Can be len(chatext) == 0 + chatdir = "./" + elif chatext is None: # Can be len(chatext) == 0 raise ValueError("Chatlog file extension is invalid") self.logger = logger self.chatdir = chatdir self.chatext = chatext self.admin = admin - self.freq_increment = freq_increment + 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, log, gen): + def store(self, tag, data, gen): chat_folder = self.chat_folder(tag=tag) chat_card = self.chat_file(tag=tag, file="card", ext=".txt") + if self.read_only: return try: if not os.path.exists(chat_folder): os.makedirs(chat_folder, exist_ok=True) self.logger.info("Storing a new chat. Folder {} created.".format(chat_folder)) - except: + except Exception: self.logger.error("Failed creating {} folder.".format(chat_folder)) return file = open(chat_card, 'w') - file.write(log) + file.write(data) file.close() + if gen is not None: chat_record = self.chat_file(tag=tag, file="record", ext=self.chatext) file = open(chat_record, 'w') file.write(gen) file.close() - def get_reader(self, filename): + def load_vocab(self, tag): + filepath = self.chat_file(tag=tag, file="record", ext=self.chatext) + try: + file = open(filepath, 'r') + record = file.read() + file.close() + return record + except Exception: + self.logger.error("Vocabulary file {} not found.".format(filepath)) + return None + + def load_reader(self, tag): + filepath = self.chat_file(tag=tag, file="card", ext=".txt") + try: + reader_file = open(filepath, 'r') + reader = reader_file.read() + reader_file.close() + return reader + except OSError: + self.logger.error("Metadata file {} not found.".format(filepath)) + return None + + def get_reader(self, tag): + reader = self.load_reader(tag) + if reader: + vocab_dump = self.load_vocab(tag) + if vocab_dump: + vocab = Generator.loads(vocab_dump) + else: + vocab = Generator() + return Reader.FromCard(reader, vocab, self.max_period, self.logger) + else: + return None + + def load_reader_old(self, filename): file = open(self.chatdir + filename, 'rb') - scribe = None + reader = None try: reader, vocab = Reader.FromFile(pickle.load(file), self) self.logger.info("Unpickled {}{}".format(self.chatdir, filename)) @@ -63,7 +99,7 @@ class Archivist(object): file.close() file = open(self.chatdir + filename, 'r') try: - scribe = Scribe.Recall(file.read(), self) + scribe = Reader.FromFile(file.read(), self) self.logger.info("Read {}{} text file".format(self.chatdir, filename)) except Exception as e: self.logger.error("Failed reading {}{}".format(self.chatdir, filename)) @@ -72,22 +108,14 @@ class Archivist(object): file.close() return scribe - def load_reader(self, filepath): - file = open(filepath.format(filename="card", ext=".txt"), 'r') - card = file.read() - file.close() - return Reader.FromCard(card, self) - - def wakeParrot(self, tag): - filepath = self.chat_file(tag=tag, file="record", ext=self.chatext) - try: - file = open(filepath, 'r') - record = file.read() - file.close() - return Generator.loads(record) - except: - self.logger.error("Record file {} not found.".format(filepath)) - return None + def chat_count(self): + count = 0 + directory = os.fsencode(self.chatdir) + for subdir in os.scandir(directory): + dirname = subdir.name.decode("utf-8") + if dirname.startswith("chat_"): + count += 1 + return count def readers_pass(self): directory = os.fsencode(self.chatdir) @@ -96,45 +124,19 @@ class Archivist(object): if dirname.startswith("chat_"): cid = dirname[5:] try: - filepath = self.chatdir + dirname + "/{filename}{ext}" - reader = self.load_reader(filepath) - self.logger.info("Chat {} contents:\n".format(cid) + reader.card.dumps()) - if self.bypass: - reader.set_period(random.randint(self.max_period//2, self.max_period)) - elif scriptorium[cid].freq() > self.max_period: - scriptorium[cid].setFreq(self.max_period) + reader = self.load_reader(cid) + # self.logger.info("Chat {} contents:\n{}".format(cid, reader.card.dumps())) + self.logger.info("Successfully read {} ({}) 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)) + elif reader.period() > self.max_period: + reader.set_period(self.max_period) + yield reader except Exception as e: self.logger.error("Failed reading {}".format(dirname)) self.logger.exception(e) raise e - """ - def wake_old(self): - scriptorium = {} - - directory = os.fsencode(self.chatdir) - for file in os.listdir(directory): - filename = os.fsdecode(file) - if filename.endswith(self.chatext): - cid = filename[:-(len(self.chatext))] - if self.filter_cids is not None: - #self.logger.info("CID " + cid) - if not cid in self.filter_cids: - continue - scriptorium[cid] = self.recall(filename) - scribe = scriptorium[cid] - if scribe is not None: - if self.bypass: - scribe.setFreq(random.randint(self.max_period//2, self.max_period)) - elif scribe.freq() > self.max_period: - scribe.setFreq(self.max_period) - self.logger.info("Loaded chat " + scribe.title() + " [" + scribe.cid() + "]" - "\n" + "\n".join(scribe.chat.dumps())) - else: - continue - return scriptorium - """ - def update(self, oldext=None): failed = [] remove = False diff --git a/generator.py b/generator.py index 17e5d45..f9deaa1 100644 --- a/generator.py +++ b/generator.py @@ -59,7 +59,7 @@ class Generator(object): # This is to mark when we want to create a Generator object from Chat data (WIP) HEAD = "\n^MESSAGE_SEPARATOR^" - TAIL = "^MESSAGE_SEPARATOR^" + TAIL = " ^MESSAGE_SEPARATOR^" def __init__(self, load=None, mode=None): if mode is not None: @@ -95,9 +95,9 @@ class Generator(object): # with the HEAD that marks the beginning of a new message and # following it with the TAIL that marks the end words = [Generator.HEAD] - text = text + " " + Generator.TAIL - words.extend(text.split()) - self.database(rewrite(text)) + text = rewrite(text + Generator.TAIL) + words.extend(text) + self.database(words) def database(self, words): # This takes a list of words and stores it in the cache, adding diff --git a/log.txt b/log.txt new file mode 100644 index 0000000..9947608 --- /dev/null +++ b/log.txt @@ -0,0 +1 @@ +Velascobot turning on. diff --git a/chatcard.py b/metadata.py similarity index 90% rename from chatcard.py rename to metadata.py index 4af559f..2a89ed7 100644 --- a/chatcard.py +++ b/metadata.py @@ -2,7 +2,7 @@ def parse_card_line(line): # This reads a line in the format 'VARIABLE=value' and gives me the value. - # See ChatCard.loadl(...) for more details + # See Metadata.loadl(...) for more details s = line.split('=', 1) if len(s) < 2: return "" @@ -10,7 +10,10 @@ def parse_card_line(line): return s[1] -class ChatCard(object): +class Metadata(object): + # This is a chat's Metadata, holding different configuration values for + # Velasco and other miscellaneous information about the chat + def __init__(self, cid, ctype, title, count=0, period=None, answer=0.5, restricted=False, silenced=False): self.id = str(cid) # The Telegram chat's ID @@ -67,7 +70,7 @@ class ChatCard(object): def loads(text): lines = text.splitlines() - return ChatCard.loadl(lines) + return Metadata.loadl(lines) def loadl(lines): # In a perfect world, I would get both the variable name and its corresponding value @@ -77,7 +80,7 @@ class ChatCard(object): version = parse_card_line(lines[0]).strip() version = version if len(version.strip()) > 1 else (lines[4] if len(lines) > 4 else "LOG_ZERO") if version == "v4" or version == "v5": - return ChatCard(cid=parse_card_line(lines[1]), + return Metadata(cid=parse_card_line(lines[1]), ctype=parse_card_line(lines[2]), title=parse_card_line(lines[3]), count=int(parse_card_line(lines[4])), @@ -87,7 +90,7 @@ class ChatCard(object): silenced=(parse_card_line(lines[8]) == 'True') ) elif version == "v3": - return ChatCard(cid=parse_card_line(lines[1]), + return Metadata(cid=parse_card_line(lines[1]), ctype=parse_card_line(lines[2]), title=parse_card_line(lines[3]), count=int(parse_card_line(lines[7])), @@ -96,7 +99,7 @@ class ChatCard(object): restricted=(parse_card_line(lines[6]) == 'True') ) elif version == "v2": - return ChatCard(cid=parse_card_line(lines[1]), + return Metadata(cid=parse_card_line(lines[1]), ctype=parse_card_line(lines[2]), title=parse_card_line(lines[3]), count=int(parse_card_line(lines[6])), @@ -107,7 +110,7 @@ class ChatCard(object): # At some point I decided to number the versions of each dictionary format, # but this was not always the case. This is what you get if you try to read # whatever there is in very old files where the version should be - return ChatCard(cid=lines[0], + return Metadata(cid=lines[0], ctype=lines[1], title=lines[2], count=int(lines[5]), @@ -115,7 +118,7 @@ class ChatCard(object): ) else: # This is for the oldest of files - return ChatCard(cid=lines[0], + return Metadata(cid=lines[0], ctype=lines[1], title=lines[2], period=int(lines[3]) diff --git a/chatreader.py b/reader.py similarity index 55% rename from chatreader.py rename to reader.py index beb486c..4189fa9 100644 --- a/chatreader.py +++ b/reader.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import random -from chatcard import ChatCard, parse_card_line +from metadata import Metadata, parse_card_line from generator import Generator @@ -25,117 +25,121 @@ class Memory(object): self.content = content -class ChatReader(object): +class Reader(object): + # This is a chat Reader object, in charge of managing the parsing of messages + # for a specific chat, and holding said chat's metadata + TAG_PREFIX = "^IS_" STICKER_TAG = "^IS_STICKER^" ANIM_TAG = "^IS_ANIMATION^" VIDEO_TAG = "^IS_VIDEO^" - def __init__(self, chatcard, max_period, logger): - self.card = chatcard + def __init__(self, metadata, vocab, max_period, logger): + self.meta = metadata + self.vocab = vocab self.max_period = max_period self.short_term_mem = [] - self.countdown = self.card.period + self.countdown = self.meta.period self.logger = logger - def FromChat(chat, max_period, logger, newchat=False): - # Create a new ChatReader from a Chat object - card = ChatCard(chat.id, chat.type, get_chat_title(chat)) - return ChatReader(card, max_period, logger) + def FromChat(chat, max_period, logger): + # Create a new Reader from a Chat object + meta = Metadata(chat.id, chat.type, get_chat_title(chat)) + vocab = Generator() + return Reader(meta, vocab, max_period, logger) - def FromData(data, max_period, logger): - # Create a new ChatReader from a whole Chat history (WIP) + def FromHistory(history, vocab, max_period, logger): + # Create a new Reader from a whole Chat history (WIP) return None - def FromCard(card, max_period, logger): - # Create a new ChatReader from a card's file dump - chatcard = ChatCard.loads(card) - return ChatReader(chatcard, max_period, logger) + def FromCard(meta, vocab, max_period, logger): + # Create a new Reader from a meta's file dump + metadata = Metadata.loads(meta) + return Reader(metadata, vocab, max_period, logger) - def FromFile(text, max_period, logger): - # Load a ChatReader from a file's text string + def FromFile(text, max_period, logger, vocab=None): + # Load a Reader from a file's text string (obsolete) lines = text.splitlines() version = parse_card_line(lines[0]).strip() version = version if len(version.strip()) > 1 else lines[4] logger.info("Dictionary version: {} ({} lines)".format(version, len(lines))) - vocab = None if version == "v4" or version == "v5": - return ChatReader.FromCard(text, max_period, logger) + return Reader.FromCard(text, vocab, max_period, logger) # I stopped saving the chat metadata and the cache together elif version == "v3": - card = ChatCard.loadl(lines[0:8]) + meta = Metadata.loadl(lines[0:8]) cache = '\n'.join(lines[9:]) vocab = Generator.loads(cache) elif version == "v2": - card = ChatCard.loadl(lines[0:7]) + meta = Metadata.loadl(lines[0:7]) cache = '\n'.join(lines[8:]) vocab = Generator.loads(cache) elif version == "dict:": - card = ChatCard.loadl(lines[0:6]) + meta = Metadata.loadl(lines[0:6]) cache = '\n'.join(lines[6:]) vocab = Generator.loads(cache) else: - card = ChatCard.loadl(lines[0:4]) + meta = Metadata.loadl(lines[0:4]) cache = lines[4:] vocab = Generator(load=cache, mode=Generator.MODE_LIST) - # raise SyntaxError("ChatReader: ChatCard format unrecognized.") - s = ChatReader(card, max_period, logger) - return (s, vocab) + # raise SyntaxError("Reader: Metadata format unrecognized.") + r = Reader(meta, vocab, max_period, logger) + return r - def archive(self, vocab): + def archive(self): # Returns a nice lice little tuple package for the archivist to save to file. # Also commits to long term memory any pending short term memories - self.commit_long_term(vocab) - return (self.card.id, self.card.dumps(), vocab) + self.commit_memory() + return (self.meta.id, self.meta.dumps(), self.vocab.dumps()) def check_type(self, t): # Checks type. Returns "True" for "group" even if it's supergroup - return t in self.card.type + return t in self.meta.type def exactly_type(self, t): # Hard check - return t == self.card.type + return t == self.meta.type def set_title(self, title): - self.card.title = title + self.meta.title = title def set_period(self, period): if period < self.countdown: self.countdown = max(period, 1) - return self.card.set_period(min(period, self.max_period)) + return self.meta.set_period(min(period, self.max_period)) def set_answer(self, prob): - return self.card.set_answer(prob) + return self.meta.set_answer(prob) def cid(self): - return str(self.card.id) + return str(self.meta.id) def count(self): - return self.card.count + return self.meta.count def period(self): - return self.card.period + return self.meta.period def title(self): - return self.card.title + return self.meta.title def answer(self): - return self.card.answer + return self.meta.answer def ctype(self): - return self.card.type + return self.meta.type def is_restricted(self): - return self.card.restricted + return self.meta.restricted def toggle_restrict(self): - self.card.restricted = (not self.card.restricted) + self.meta.restricted = (not self.meta.restricted) def is_silenced(self): - return self.card.silenced + return self.meta.silenced def toggle_silence(self): - self.card.silenced = (not self.card.silenced) + self.meta.silenced = (not self.meta.silenced) def is_answering(self): rand = random.random() @@ -151,24 +155,26 @@ class ChatReader(object): self.short_term_mem.append(mem) def random_memory(self): + if len(self.short_term_mem) == 0: + return None mem = random.choice(self.short_term_mem) return mem.id def reset_countdown(self): - self.countdown = self.card.period + self.countdown = self.meta.period def read(self, message): mid = str(message.message_id) if message.text is not None: - self.read(mid, message.text) + self.learn(mid, message.text) elif message.sticker is not None: - self.learn_drawing(mid, ChatReader.STICKER_TAG, message.sticker.file_id) + self.learn_drawing(mid, Reader.STICKER_TAG, message.sticker.file_id) elif message.animation is not None: - self.learn_drawing(mid, ChatReader.ANIM_TAG, message.animation.file_id) + self.learn_drawing(mid, Reader.ANIM_TAG, message.animation.file_id) elif message.video is not None: - self.learn_drawing(mid, ChatReader.VIDEO_TAG, message.video.file_id) - self.card.count += 1 + self.learn_drawing(mid, Reader.VIDEO_TAG, message.video.file_id) + self.meta.count += 1 def learn_drawing(self, mid, tag, drawing): self.learn(mid, tag + " " + drawing) @@ -178,13 +184,10 @@ class ChatReader(object): return self.add_memory(mid, text) - def commit_long_term(self, vocab): + def commit_memory(self): for mem in self.short_term_mem: - vocab.add(mem.content) + self.vocab.add(mem.content) self.short_term_mem = [] - """ - def learnFrom(self, scribe): - self.card.count += scribe.chat.count - self.vocab.cross(scribe.vocab) - """ + def generate_message(self, max_len): + return self.vocab.generate(size=max_len, silence=self.is_silenced()) diff --git a/speaker.py b/speaker.py index f07788c..d7e4391 100644 --- a/speaker.py +++ b/speaker.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 import random -from chatreader import ChatReader as Reader -from telegram.error import * +from reader import Reader, get_chat_title +from telegram.error import TimedOut def send(bot, cid, text, replying=None, formatting=None, logger=None, **kwargs): @@ -33,105 +33,117 @@ class Speaker(object): ModeFixed = "FIXED_MODE" ModeChance = "MODE_CHANCE" - def __init__(self, name, username, archivist, logger, + def __init__(self, username, archivist, logger, nicknames=[], reply=0.1, repeat=0.05, wakeup=False, mode=ModeFixed ): - self.name = name + self.names = nicknames self.username = username self.archivist = archivist - self.scriptorium = archivist.wakeScriptorium() logger.info("----") logger.info("Finished loading.") - logger.info("Loaded {} chats.".format(len(self.scriptorium))) + logger.info("Loaded {} chats.".format(archivist.chat_count())) logger.info("----") self.wakeup = wakeup self.logger = logger self.reply = reply self.repeat = repeat - self.filterCids = archivist.filterCids + self.filter_cids = archivist.filter_cids self.bypass = archivist.bypass + self.current_reader = None - def announce(self, announcement, check=(lambda _: True)): - for scribe in self.scriptorium: + 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(): try: - if check(scribe): - send(bot, scribe.cid(), announcement) - logger.info("Waking up on chat {}".format(scribe.cid())) - except: + if check(reader): + send(bot, reader.cid(), announcement) + self.logger.info("Sending announcement to chat {}".format(reader.cid())) + except Exception: pass def wake(self, bot, wake): + # Sends a wake-up message as announcement to all chats that + # are groups if self.wakeup: - def check(scribe): - return scribe.checkType("group") - self.announce(wake, check) + def group_check(reader): + return reader.check_type("group") + self.announce(bot, wake, group_check) - def getScribe(self, chat): + def load_reader(self, chat): cid = str(chat.id) - if not cid in self.scriptorium: - scribe = Reader.FromChat(chat, self.archivist, newchat=True) - self.scriptorium[cid] = scribe - return scribe - else: - return self.scriptorium[cid] + if self.current_reader is not None and cid == self.current_reader.cid(): + return - def shouldReply(self, message, scribe): - if not self.bypass and scribe.isRestricted(): + if self.current_reader is not None: + self.current_reader.commit_memory() + self.save() + + reader = self.archivist.get_reader(cid) + if not reader: + reader = Reader.FromChat(chat, self.archivist.max_period, self.logger) + self.current_reader = reader + + def get_reader(self, cid): + if self.current_reader is None or cid != self.current_reader.cid(): + return self.archivist.get_reader(cid) + + return self.current_reader + + def mentioned(self, text): + if self.username in text: + return True + for name in self.names: + if name in text and "@{}".format(name) not in text: + return True + return False + + def should_reply(self, message): + if not self.bypass and self.current_reader.is_restricted(): user = message.chat.get_member(message.from_user.id) - if not self.userIsAdmin(user): + if not self.user_is_admin(user): # update.message.reply_text("You do not have permissions to do that.") return False replied = message.reply_to_message text = message.text.casefold() if message.text else "" - return ( ((replied is not None) and (replied.from_user.name == self.username)) or - (self.username in text) or - (self.name in text and "@{}".format(self.name) not in text) - ) + return (((replied is not None) and (replied.from_user.name == self.username)) + or (self.mentioned(text))) - def store(self, scribe): - if self.parrot is None: - raise ValueError("Tried to store a Parrot that is None.") + def save(self): + if self.current_reader is None: + raise ValueError("Tried to store a None Reader.") else: - scribe.store(self.parrot.dumps()) + self.archivist.store(*self.current_reader.archive()) - def loadParrot(self, scribe): - newParrot = False - self.parrot = self.archivist.wakeParrot(scribe.cid()) - if self.parrot is None: - newParrot = True - self.parrot = Markov() - scribe.teachParrot(self.parrot) - self.store(scribe) - return newParrot - - def read(self, bot, update): + def read(self, update, context): + if update.message is None: + return chat = update.message.chat - scribe = self.getScribe(chat) - scribe.learn(update.message) + self.load_reader(chat) + self.current_reader.read(update.message) - if self.shouldReply(update.message, scribe) and scribe.isAnswering(): - self.say(bot, scribe, replying=update.message.message_id) + if self.should_reply(update.message) and self.current_reader.is_answering(): + self.say(context.bot, replying=update.message.message_id) return - title = getTitle(update.message.chat) - if title != scribe.title(): - scribe.setTitle(title) + title = get_chat_title(update.message.chat) + if title != self.current_reader.title(): + self.current_reader.set_title(title) - scribe.countdown -= 1 - if scribe.countdown < 0: - scribe.resetCountdown() - rid = scribe.getReference() if random.random() <= self.reply else None - self.say(bot, scribe, replying=rid) - elif (scribe.freq() - scribe.countdown) % self.archivist.saveCount == 0: - self.loadParrot(scribe) + self.current_reader.countdown -= 1 + if self.current_reader.countdown < 0: + self.current_reader.reset_countdown() + rid = self.current_reader.random_memory() if random.random() <= self.reply else None + self.say(context.bot, replying=rid) + elif (self.current_reader.period() - self.current_reader.countdown) % self.archivist.save_count == 0: + self.save() - def speak(self, bot, update): + def speak(self, update, context): chat = (update.message.chat) - scribe = self.getScribe(chat) + self.load_reader(chat) - if not self.bypass and scribe.isRestricted(): + if not self.bypass and self.current_reader.is_restricted(): user = update.message.chat.get_member(update.message.from_user.id) - if not self.userIsAdmin(user): + if not self.user_is_admin(user): # update.message.reply_text("You do not have permissions to do that.") return @@ -140,148 +152,153 @@ class Speaker(object): rid = replied.message_id if replied else mid words = update.message.text.split() if len(words) > 1: - scribe.learn(' '.join(words[1:])) - self.say(bot, scribe, replying=rid) + self.current_reader.read(' '.join(words[1:])) + self.say(context.bot, replying=rid) - def userIsAdmin(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("Bot Creator ID is {}".format(str(self.archivist.admin))) - return ((member.status == 'creator') or - (member.status == 'administrator') or - (member.user.id == self.archivist.admin)) + return ((member.status == 'creator') + or (member.status == 'administrator') + or (member.user.id == self.archivist.admin)) - def speech(self, scribe): - return self.parrot.generate_markov_text(size=self.archivist.maxLen, silence=scribe.isSilenced()) + def speech(self): + return self.current_reader.generate_message(self.archivist.max_len) - def say(self, bot, scribe, replying=None, **kwargs): - if self.filterCids is not None and not scribe.cid() in self.filterCids: + def say(self, bot, replying=None, **kwargs): + cid = self.current_reader.cid() + if self.filter_cids is not None and cid not in self.filter_cids: return - self.loadParrot(scribe) try: - send(bot, scribe.cid(), self.speech(scribe), replying, logger=self.logger, **kwargs) + send(bot, cid, self.speech(), replying, logger=self.logger, **kwargs) if self.bypass: - maxFreq = self.archivist.maxFreq - scribe.setFreq(random.randint(maxFreq//4, maxFreq)) + max_period = self.archivist.max_period + self.current_reader.set_period(random.randint(max_period // 4, max_period)) if random.random() <= self.repeat: - send(bot, scribe.cid(), self.speech(scribe), logger=self.logger, **kwargs) + send(bot, cid, self.speech(), logger=self.logger, **kwargs) except TimedOut: - scribe.setFreq(scribe.freq() + self.archivist.freqIncrement) - self.logger.warning("Increased period for chat {} [{}]".format(scribe.title(), scribe.cid())) + self.current_reader.set_period(self.current_reader.period() + self.archivist.period_inc) + self.logger.warning("Increased period for chat {} [{}]".format(self.current_reader.title(), cid)) except Exception as e: self.logger.error("Sending a message caused error:") - self.logger.error(e) + raise e - def getCount(self, bot, update): + def get_count(self, update, context): cid = str(update.message.chat.id) - scribe = self.scriptorium[cid] - num = str(scribe.count()) if self.scriptorium[cid] else "no" + reader = self.get_reader(cid) + + num = str(reader.count()) if reader else "no" update.message.reply_text("I remember {} messages.".format(num)) - def getChats(self, bot, update): - lines = ["[{}]: {}".format(cid, self.scriptorium[cid].title()) for cid in self.scriptorium] - list = "\n".join(lines) - update.message.reply_text( "\n\n".join(["I have the following chats:", list]) ) + def get_chats(self, update, context): + lines = ["[{}]: {}".format(reader.cid(), reader.title()) for reader in self.archivist.readers_pass] + chat_list = "\n".join(lines) + update.message.reply_text("I have the following chats:\n\n" + chat_list) - def freq(self, bot, update): + def period(self, update, context): chat = update.message.chat - scribe = self.getScribe(chat) + reader = self.get_reader(str(chat.id)) words = update.message.text.split() if len(words) <= 1: - update.message.reply_text("The current speech period is {}".format(scribe.freq())) + update.message.reply_text("The current speech period is {}".format(reader.period())) return - if scribe.isRestricted(): + if reader.is_restricted(): user = update.message.chat.get_member(update.message.from_user.id) - if not self.userIsAdmin(user): + if not self.user_is_admin(user): update.message.reply_text("You do not have permissions to do that.") return try: - freq = int(words[1]) - freq = scribe.setFreq(freq) - update.message.reply_text("Period of speaking set to {}.".format(freq)) - scribe.store(None) - except: - update.message.reply_text("Format was confusing; period unchanged from {}.".format(scribe.freq())) + 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()) + except Exception: + update.message.reply_text("Format was confusing; period unchanged from {}.".format(reader.period())) - def answer(self, bot, update): + def answer(self, update, context): chat = update.message.chat - scribe = self.getScribe(chat) + reader = self.get_reader(str(chat.id)) words = update.message.text.split() if len(words) <= 1: - update.message.reply_text("The current answer probability is {}".format(scribe.answer())) + update.message.reply_text("The current answer probability is {}".format(reader.answer())) return - if scribe.isRestricted(): + if reader.is_restricted(): user = update.message.chat.get_member(update.message.from_user.id) - if not self.userIsAdmin(user): + if not self.user_is_admin(user): update.message.reply_text("You do not have permissions to do that.") return try: - answ = float(words[1]) - answ = scribe.setAnswer(answ) - update.message.reply_text("Answer probability set to {}.".format(answ)) - scribe.store(None) - except: - update.message.reply_text("Format was confusing; answer probability unchanged from {}.".format(scribe.answer())) + answer = float(words[1]) + answer = reader.set_answer(answer) + update.message.reply_text("Answer probability set to {}.".format(answer)) + self.archivist.store(*reader.archive()) + except Exception: + update.message.reply_text("Format was confusing; answer probability unchanged from {}.".format(reader.answer())) - def restrict(self, bot, update): + def restrict(self, update, context): if "group" not in update.message.chat.type: update.message.reply_text("That only works in groups.") return chat = update.message.chat user = chat.get_member(update.message.from_user.id) - scribe = self.getScribe(chat) - if scribe.isRestricted(): - if not self.userIsAdmin(user): + reader = self.get_reader(str(chat.id)) + + if reader.is_restricted(): + if not self.user_is_admin(user): update.message.reply_text("You do not have permissions to do that.") return - scribe.restrict() - allowed = "let only admins" if scribe.isRestricted() else "let everyone" + 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()) - def silence(self, bot, update): + def silence(self, update, context): if "group" not in update.message.chat.type: update.message.reply_text("That only works in groups.") return chat = update.message.chat user = chat.get_member(update.message.from_user.id) - scribe = self.getScribe(chat) - if scribe.isRestricted(): - if not self.userIsAdmin(user): + reader = self.get_reader(str(chat.id)) + + if reader.is_restricted(): + if not self.user_is_admin(user): update.message.reply_text("You do not have permissions to do that.") return - scribe.silence() - allowed = "avoid mentioning" if scribe.isSilenced() else "mention" + 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()) - def who(self, bot, update): + def who(self, update, context): msg = update.message usr = msg.from_user cht = msg.chat chtname = cht.title if cht.title else cht.first_name + rdr = self.get_reader(str(cht.id)) answer = ("You're **{name}**, with username `{username}`, and " "id `{uid}`.\nYou're messaging in the chat named __{cname}__," " of type {ctype}, with id `{cid}`, and timestamp `{tstamp}`." ).format(name=usr.full_name, username=usr.username, uid=usr.id, cname=chtname, cid=cht.id, - ctype=scribe.type(), tstamp=str(msg.date)) + ctype=rdr.ctype(), tstamp=str(msg.date)) msg.reply_markdown(answer) - def where(self, bot, update): - print("THEY'RE ASKING WHERE") + def where(self, update, context): msg = update.message chat = msg.chat - scribe = self.getScribe(chat) - if scribe.isRestricted() and scribe.isSilenced(): + reader = self.get_reader(str(chat.id)) + if reader.is_restricted() and reader.is_silenced(): permissions = "restricted and silenced" - elif scribe.isRestricted(): + elif reader.is_restricted(): permissions = "restricted but not silenced" - elif scribe.isSilenced(): + elif reader.is_silenced(): permissions = "not restricted but silenced" else: permissions = "neither restricted nor silenced" @@ -289,8 +306,8 @@ class Speaker(object): answer = ("You're messaging in the chat of saved title __{cname}__," " with id `{cid}`, message count {c}, period {p}, and answer " "probability {a}.\n\nThis chat is {perm}." - ).format(cname=scribe.title(), cid=scribe.cid(), - c=scribe.count(), p=scribe.freq(), a=scribe.answer(), - perm=permissions) + ).format(cname=reader.title(), cid=reader.cid(), + c=reader.count(), p=reader.period(), + a=reader.answer(), perm=permissions) msg.reply_markdown(answer) diff --git a/velasco.py b/velasco.py index 162b748..e854d28 100644 --- a/velasco.py +++ b/velasco.py @@ -1,8 +1,9 @@ #!/usr/bin/env python3 -import logging, argparse +import logging +import argparse from telegram.ext import Updater, CommandHandler, MessageHandler, Filters -from telegram.error import * +# from telegram.error import * from archivist import Archivist from speaker import Speaker @@ -38,7 +39,7 @@ help_msg = """I answer to the following commands: /explain - I explain how I work. /help - I send this message. /count - I tell you how many messages from this chat I remember. -/freq - Change the frequency of my messages. (Maximum of 100000) +/period - Change the period of my messages. (Maximum of 100000) /speak - Forces me to speak. /answer - Change the probability to answer to a reply. (Decimal between 0 and 1). /restrict - Toggle restriction of configuration commands to admins only. @@ -47,7 +48,7 @@ help_msg = """I answer to the following commands: about_msg = "I am yet another Markov Bot experiment. I read everything you type to me and then spit back nonsensical messages that look like yours.\n\nYou can send /explain if you want further explanation." -explanation = "I decompose every message I read in groups of 3 consecutive words, so for each consecutive pair I save the word that can follow them. I then use this to make my own messages. At first I will only repeat your messages because for each 2 words I will have very few possible following words.\n\nI also separate my vocabulary by chats, so anything I learn in one chat I will only say in that chat. For privacy, you know. Also, I save my vocabulary in the form of a json dictionary, so no logs are kept.\n\nMy default frequency in private chats is one message of mine from each 2 messages received, and in group chats it\'s 10 messages I read for each message I send." +explanation = "I decompose every message I read in groups of 3 consecutive words, so for each consecutive pair I save the word that can follow them. I then use this to make my own messages. At first I will only repeat your messages because for each 2 words I will have very few possible following words.\n\nI also separate my vocabulary by chats, so anything I learn in one chat I will only say in that chat. For privacy, you know. Also, I save my vocabulary in the form of a json dictionary, so no logs are kept.\n\nMy default period in private chats is one message of mine from each 2 messages received, and in group chats it\'s 10 messages I read for each message I send." def static_reply(text, format=None): @@ -56,15 +57,16 @@ def static_reply(text, format=None): return reply -def error(bot, update, error): - logger.warning('Update "{}" caused error "{}"'.format(update, error)) +def error(update, context): + logger.warning('The following update:\n"{}"\n\nCaused the following error:\n"{}"'.format(update, context.error)) + # raise error def stop(bot, update): - scribe = speakerbot.getScribe(update.message.chat.id) - #del chatlogs[chatlog.id] - #os.remove(LOG_DIR + chatlog.id + LOG_EXT) - logger.warning("I got blocked by user {} [{}]".format(scribe.title(), scribe.cid())) + reader = speakerbot.get_reader(str(update.message.chat.id)) + # del chatlogs[chatlog.id] + # os.remove(LOG_DIR + chatlog.id + LOG_EXT) + logger.warning("I got blocked by user {} [{}]".format(reader.title(), reader.cid())) def main(): @@ -73,38 +75,42 @@ def main(): 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.') args = parser.parse_args() # Create the EventHandler and pass it your bot's token. - updater = Updater(args.token) + updater = Updater(args.token, use_context=True) - #filterCids=["-1001036575277", "-1001040087584", str(args.admin_id)] - filterCids = None + filter_cids = args.filter + if filter_cids: + filter_cids.append(str(args.admin_id)) archivist = Archivist(logger, chatdir="chatlogs/", chatext=".vls", admin=args.admin_id, - filterCids=filterCids, - readOnly=False + filter_cids=filter_cids, + read_only=False ) - speakerbot = Speaker("velasco", "@" + username, archivist, logger, wakeup=args.wakeup) + username = updater.bot.get_me().username + speakerbot = Speaker("@" + username, archivist, logger, nicknames=args.nicknames, wakeup=args.wakeup) # Get the dispatcher to register handlers dp = updater.dispatcher # on different commands - answer in Telegram - dp.add_handler(CommandHandler("start", static_reply(start_msg) )) - dp.add_handler(CommandHandler("about", static_reply(about_msg) )) - dp.add_handler(CommandHandler("explain", static_reply(explanation) )) - dp.add_handler(CommandHandler("help", static_reply(help_msg) )) - dp.add_handler(CommandHandler("count", speakerbot.getCount)) - dp.add_handler(CommandHandler("period", speakerbot.freq)) - dp.add_handler(CommandHandler("list", speakerbot.getChats, 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("start", static_reply(start_msg))) + dp.add_handler(CommandHandler("about", static_reply(about_msg))) + dp.add_handler(CommandHandler("explain", static_reply(explanation))) + 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("user", get_name, Filters.chat(chat_id=archivist.admin))) + # dp.add_handler(CommandHandler("id", get_id)) dp.add_handler(CommandHandler("stop", stop)) dp.add_handler(CommandHandler("speak", speakerbot.speak)) dp.add_handler(CommandHandler("answer", speakerbot.answer)) @@ -130,5 +136,6 @@ def main(): # start_polling() is non-blocking and will stop the bot gracefully. updater.idle() + if __name__ == '__main__': main()