Markov-Chain brain/bot for •AZ•BOT•

markov.py

"""
Markovbot mod

Depends on:
    core mods (>= 1.3.0)

Config options:

markovbot:
  # These are good defaults. Change with care.
  stop_word: "\n"
  chain_length: 3
  max_words: 10000
  
  # File to save the brain words to.
  # It's append only.
  brain_file: markov.txt
  
  # Name the markovbot responds to.
  prefix: "clever: "

  # Name the markovbot speaks to.
  #
  # If not present, will default to whomever
  # addressed the bot. If empty, will speak
  # to everyone.
  #talk_to:
"""

from collections import defaultdict
from pubsub import pub
import _util
import random
import exceptions

markov = defaultdict(list)
config = {}

def add_to_brain(msg):
    n = len(markov)
    if 'brain_file' in config:
        f = open(config['brain_file'], 'a+')
        f.write(msg + '\n')
        f.close()
    buf = [config['stop_word']] * config['chain_length']
    for word in msg.split():
        markov[tuple(buf)].append(word)
        del buf[0]
        buf.append(word)
    markov[tuple(buf)].append(config['stop_word'])
    if len(markov) > n:
        print "Added %d words to brain." % (len(markov) - n)

def generate_sentence(msg):
    if len(markov) < 50:
        f = open(config['brain_file'], 'r')
        brain = f.read()
        f.close()
        print "Brain loading up..."
        for phrase in brain.split(config['stop_word']):
          add_to_brain(phrase)
    buf = msg.split()[:config['chain_length']]
    if len(msg.split()) > config['chain_length']:
        message = buf[:]
    else:
        message = []
        for i in xrange(config['chain_length']):
            try:
                message.append(random.choice(markov[random.choice(markov.keys())]))
            except IndexError:
                pass
    
    for i in xrange(config['max_words']):
        try:
            next_word = random.choice(markov[tuple(buf)])
        except IndexError:
            continue
        
        if next_word == config['stop_word']:
            break
        
        message.append(next_word)
        del buf[0]
        buf.append(next_word)
    return ' '.join(message)


def scanner(bot, user, channel, message, topic=pub.AUTO_TOPIC):
    coords = {
        'bot': bot,
        'user': user,
        'channel': channel
    }
    
    if 'markovbot' in bot.config:
        global config
        config = bot.config['markovbot']
    else:
        print "[!] Markov: not configured."
        return None
    
    if bot.network['name'] == topic.getNameTuple()[-1]:
        if message.startswith(config['prefix']):
            line = message[len(config['prefix']):]
            
            add_to_brain(line)
            gend = ""
            while len(gend) < 1:
              gend = generate_sentence(line).strip().replace('\n','')
            
            if 'talk_to' in config:
                prep = (config['talk_to'], gend)
            else:
                prep = (user, gend)
            
            _util.answer(coords, "%s: %s" % prep)



pub.subscribe(scanner, 'irc.message')