diff options
Diffstat (limited to 'plugins/ai.py')
-rw-r--r-- | plugins/ai.py | 177 |
1 files changed, 177 insertions, 0 deletions
diff --git a/plugins/ai.py b/plugins/ai.py new file mode 100644 index 00000000..d86b2393 --- /dev/null +++ b/plugins/ai.py @@ -0,0 +1,177 @@ +""" +Recreational plugin. + +Message parser that can generate sentences based on what he has already seen +before. + +""" +from plugin import BasePlugin +from random import choice +from re import split as rsplit +import pickle +import tabs + +class Dico(object): + def __init__(self): + self.start_words = [] + self.end_words = [] + self.words = {} + + def add_next(self, word, next): + w = self.words[word] + if next in w: + w[next] += 1 + else: + w[next] = 1 + + def add_word(self, word): + if not word in self.words: + self.words[word] = {} + + def select_next(self, word): + d = sorted(self.words[word], key=lambda w: self.words[word][w], reverse=True) + if not d: + return '' + nexts = d[:10] + for i in range(0, len(d) // 10): + nexts.append(choice(d[9:])) + return choice(nexts) + + def create_sentence(self, length): + if not self.start_words: + return '' + current = choice(self.start_words) + i = 1 + sent = current.capitalize() + while current and self.words[current] and i < length: + current = self.select_next(current) + sent += " " + current + i += 1 + return sent + + def save(self, fname): + file = open(fname, 'wb') + pickle.dump(self, file) + file.close + +spaces = ' ' +end_sentence = ['.', '?', '!'] + +def end_re(): + r = '(' + for i in end_sentence[:]: + end_sentence.append('%s ' % i) + i = '\%s'% i + r += '%s$|%s |' % (i, i) + r = r[:-1] + r += ')' + return r + +end = end_re() + + +class Analyzer(object): + dico = None + def __init__(self): + pass + + def parse(self, text): + text = text.replace('\n', '') + res = rsplit(end, text) + for i in res[:]: + if i == '': + continue + elif i in end_sentence: + continue + self.analyze(i) + + def analyze(self, text): + prev = None + for word in rsplit('[%s]'%spaces, text): + if word in spaces: continue + word = word.lower() + self.dico.add_word(word) + if prev: + self.dico.add_next(prev, word) + else: + self.dico.start_words.append(word) + prev = word + +class Plugin(BasePlugin): + def init(self): + self.add_event_handler('groupchat_message', self.on_groupchat_message) + self.add_tab_command(tabs.MucTab, 'random', self.command_random, '/random [n]\nRandom: Send a random message, if n is provided and is integer > 1, the message will have a maximum number of n words', None) + self.add_tab_command(tabs.MucTab, 'start', self.command_start, '/start\nStart: Start parsing the messages', None) + self.add_tab_command(tabs.MucTab, 'stop', self.command_stop, '/stop\nStop: Stop parsing the messages', None) + self.add_tab_command(tabs.MucTab, 'flush', self.command_flush, '/flush\nFlush: Flush the database', None) + self.add_tab_command(tabs.MucTab, 'save', self.command_save, '/save <filepath>\nSave: Save the database to a file', None) + self.add_tab_command(tabs.MucTab, 'load', self.command_load, '/load <filepath>\nLoad: Load the database from a file', None) + self.tabs = {} + self.analyzer = Analyzer() + + def command_start(self, arg): + name = self.core.current_tab().get_name() + if not name in self.tabs: + self.tabs[name] = Dico() + self.core.information('Started analyzing in %s' % name, 'Info') + else: + self.core.information('Already started', 'Info') + + def command_stop(self, arg): + name = self.core.current_tab().get_name() + if name in self.tabs: + del self.tabs[name] + self.core.information('Stopped analyzing in %s' % name, 'Info') + else: + self.core.information('Nothing to stop', 'Info') + + def command_save(self, arg): + name = self.core.current_tab().get_name() + if name in self.tabs: + try: + self.tabs[name].save(arg) + except: + self.core.information('Could not save the file', 'Info') + else: + self.core.information('Nothing to save', 'Info') + + def command_flush(self, arg): + name = self.core.current_tab().get_name() + if name in self.tabs: + del self.tabs[name] + self.tabs[name] = Dico() + self.core.information('Database flushed', 'Info') + else: + self.core.information('Nothing to flush', 'Info') + + def command_load(self, arg): + name = self.core.current_tab().get_name() + try: + file = open(arg, 'rb') + self.tabs[name] = pickle.load(file) + file.close() + self.core.information('File loaded', 'Info') + except: + self.core.information('Could not load the file', 'Info') + + def on_groupchat_message(self, message): + if not message['body']: + return + jid = message['from'] + if jid.bare not in self.tabs or jid.resource == self.core.current_tab().own_nick: + return + jid = jid.bare + + self.analyzer.dico = self.tabs[jid] + self.analyzer.parse(message['body']) + + def command_random(self, arg): + name = self.core.current_tab().get_name() + try: + i = int(arg) + if i < 1: + i = 1 + except: + i = 25 + if name in self.tabs: + self.core.send_message(self.tabs[name].create_sentence(i)) |