diff --git a/archivist.py b/archivist.py index 29d72db..69d35fb 100644 --- a/archivist.py +++ b/archivist.py @@ -1,5 +1,5 @@ -import os, random +import os from reader import Reader from generator import Generator @@ -28,7 +28,7 @@ class Archivist(object): def chat_file(self, *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, vocab_dumper): chat_folder = self.chat_folder(tag=tag) chat_card = self.chat_file(tag=tag, file="card", ext=".txt") @@ -45,17 +45,17 @@ class Archivist(object): file.write(data) file.close() - if gen is not None: + if vocab_dumper is not None: chat_record = self.chat_file(tag=tag, file="record", ext=self.chatext) file = open(chat_record, 'w', encoding="utf-16") - file.write(gen) + vocab_dumper(file) file.close() def load_vocab(self, tag): filepath = self.chat_file(tag=tag, file="record", ext=self.chatext) try: file = open(filepath, 'r', encoding="utf-16") - record = file.read() + record = Generator.load(file) file.close() return record except Exception as e: @@ -89,10 +89,8 @@ class Archivist(object): 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 = self.load_vocab(tag) + if not vocab: vocab = Generator() return Reader.FromCard(reader, vocab, self.max_period, self.logger) else: diff --git a/generator.py b/generator.py index bd1bb50..590eead 100644 --- a/generator.py +++ b/generator.py @@ -55,6 +55,9 @@ class Generator(object): MODE_LIST = "MODE_LIST" # This is to mark when we want to create a Generator object from a given list of words + MODE_DICT = "MODE_DICT" + # This is to mark when we want to create a Generator object from a given dictionary + MODE_CHAT_DATA = "MODE_CHAT_DATA" # This is to mark when we want to create a Generator object from Chat data (WIP) @@ -69,6 +72,8 @@ class Generator(object): elif mode == Generator.MODE_LIST: self.cache = {} self.load_list(load) + elif mode == Generator.MODE_DICT: + self.cache = load else: self.cache = {} # The cache is where we store our words @@ -82,6 +87,9 @@ class Generator(object): # Dumps the cache dictionary into a JSON-formatted string return json.dumps(self.cache, ensure_ascii=False) + def dump(self, f): + json.dump(self.cache, f, ensure_ascii=False, indent='') + def loads(dump): # Loads the cache dictionary from a JSON-formatted string if len(dump) == 0: @@ -90,6 +98,9 @@ class Generator(object): # otherwise return Generator(load=dump, mode=Generator.MODE_JSON) + def load(self, f): + return Generator(load=json.load(f), mode=Generator.MODE_DICT) + def add(self, text): # This takes a string and stores it in the cache, preceding it # with the HEAD that marks the beginning of a new message and diff --git a/reader.py b/reader.py index 4189fa9..e4ce04b 100644 --- a/reader.py +++ b/reader.py @@ -90,7 +90,7 @@ class Reader(object): # 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_memory() - return (self.meta.id, self.meta.dumps(), self.vocab.dumps()) + return (self.meta.id, self.meta.dumps(), self.vocab.dump) def check_type(self, t): # Checks type. Returns "True" for "group" even if it's supergroup diff --git a/speaker.py b/speaker.py index e96257f..4f07748 100644 --- a/speaker.py +++ b/speaker.py @@ -232,12 +232,10 @@ class Speaker(object): if random.random() <= self.repeat: send(bot, cid, self.speech(reader), logger=self.logger, **kwargs) except NetworkError as e: - if '429' in e.message: - self.logger.error("Error: TooManyRequests. Going mute for {} seconds.".format(self.mute_time)) - self.mute_timer = int(time.perf_counter()) - else: - self.logger.error("Sending a message caused network error:") - self.logger.exception(e) + self.logger.error("Sending a message caused network error:") + self.logger.exception(e) + self.logger.error("Going mute for {} seconds.".format(self.mute_time)) + self.mute_timer = int(time.perf_counter()) except Exception as e: self.logger.error("Sending a message caused exception:") self.logger.exception(e) diff --git a/velasco.py b/velasco.py index 48341d2..a5b8380 100644 --- a/velasco.py +++ b/velasco.py @@ -85,6 +85,8 @@ def main(): 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('-c', '--capacity', metavar='C', type=int, default=20, + help='The memory capacity for the last C updated chats. (default: 20).') 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, @@ -113,6 +115,7 @@ def main(): filter_cids=filter_cids, nicknames=args.nicknames, wakeup=args.wakeup, + memory=args.capacity, mute_time=args.mute_time, save_time=args.save_time)