145 lines
6.2 KiB
Python
145 lines
6.2 KiB
Python
"""IRC support for Markov chain learning and text generation."""
|
|
import logging
|
|
import re
|
|
|
|
import irc.client
|
|
|
|
import markov.lib as markovlib
|
|
from ircbot.lib import Plugin, reply_destination_for_event
|
|
from ircbot.models import IrcChannel
|
|
from markov.models import MarkovContext, MarkovTarget
|
|
|
|
log = logging.getLogger('markov.ircplugin')
|
|
|
|
|
|
class Markov(Plugin):
|
|
"""Build Markov chains and reply with them."""
|
|
|
|
def start(self):
|
|
"""Set up the handlers."""
|
|
self.connection.add_global_handler('pubmsg', self.handle_chatter, -20)
|
|
self.connection.add_global_handler('privmsg', self.handle_chatter, -20)
|
|
|
|
self.connection.reactor.add_global_regex_handler(['pubmsg', 'privmsg'],
|
|
r'^!markov\s+reply(\s+min=(\d+))?(\s+max=(\d+))?(\s+(.*)$|$)',
|
|
self.handle_reply, -20)
|
|
|
|
super(Markov, self).start()
|
|
|
|
def stop(self):
|
|
"""Tear down handlers."""
|
|
self.connection.remove_global_handler('pubmsg', self.handle_chatter)
|
|
self.connection.remove_global_handler('privmsg', self.handle_chatter)
|
|
|
|
self.connection.reactor.remove_global_regex_handler(['pubmsg', 'privmsg'], self.handle_reply)
|
|
|
|
super(Markov, self).stop()
|
|
|
|
def handle_reply(self, connection, event, match):
|
|
"""Generate a reply to one line, without learning it."""
|
|
target = reply_destination_for_event(event)
|
|
|
|
min_size = 15
|
|
max_size = 30
|
|
context = self.get_or_create_target_context(target)
|
|
|
|
if match.group(2):
|
|
min_size = int(match.group(2))
|
|
if match.group(4):
|
|
max_size = int(match.group(4))
|
|
|
|
if match.group(5) != '':
|
|
line = match.group(6)
|
|
topics = [x for x in line.split(' ') if len(x) >= 3]
|
|
|
|
return self.bot.reply(event, " ".join(markovlib.generate_line(context, topics=topics,
|
|
min_words=min_size, max_words=max_size)))
|
|
else:
|
|
return self.bot.reply(event, " ".join(markovlib.generate_line(context, min_words=min_size,
|
|
max_words=max_size)))
|
|
|
|
def handle_chatter(self, connection, event):
|
|
"""Learn from IRC chatter."""
|
|
what = event.arguments[0]
|
|
who = irc.client.NickMask(event.source).nick
|
|
target = reply_destination_for_event(event)
|
|
|
|
log.debug("what: '%s', who: '%s', target: '%s'", what, who, target)
|
|
# check to see whether or not we should learn from this channel
|
|
channel = None
|
|
if irc.client.is_channel(target):
|
|
channel, c = IrcChannel.objects.get_or_create(name=target, server=connection.server_config)
|
|
|
|
if channel and not channel.markov_learn_from_channel:
|
|
log.debug("not learning from %s as i've been told to ignore it", channel)
|
|
else:
|
|
# learn the line
|
|
learning_what = what
|
|
|
|
# don't learn the speaker's nick if this came over a bridge
|
|
if channel and who == channel.discord_bridge:
|
|
learning_what = ' '.join(learning_what.split(' ')[1:])
|
|
|
|
# remove our own nick and aliases from what we learn
|
|
if connection.server_config.additional_addressed_nicks:
|
|
all_nicks = '|'.join(connection.server_config.additional_addressed_nicks.split('\n') +
|
|
[connection.get_nickname()])
|
|
else:
|
|
all_nicks = connection.get_nickname()
|
|
learning_what = re.sub(r'^(({nicks})[:,]|@({nicks})[:,]?)\s+'.format(nicks=all_nicks), '', learning_what)
|
|
|
|
recursing = getattr(event, 'recursing', False)
|
|
if not recursing:
|
|
log.debug("learning %s", learning_what)
|
|
context = self.get_or_create_target_context(target)
|
|
markovlib.learn_line(learning_what, context)
|
|
|
|
log.debug("searching '%s' for '%s'", what, all_nicks)
|
|
if re.search(all_nicks, what, re.IGNORECASE) is not None:
|
|
context = self.get_or_create_target_context(target)
|
|
|
|
addressed_pattern = r'^(({nicks})[:,]|@({nicks})[:,]?)\s+(?P<addressed_msg>.*)'.format(nicks=all_nicks)
|
|
match = re.match(addressed_pattern, what, re.IGNORECASE)
|
|
if match:
|
|
# i was addressed directly, so respond, addressing
|
|
# the speaker
|
|
topics = [x for x in match.group('addressed_msg').split(' ') if len(x) >= 3]
|
|
|
|
return self.bot.reply(event, "{0:s}: {1:s}"
|
|
"".format(who, " ".join(markovlib.generate_line(context, topics=topics))))
|
|
else:
|
|
# i wasn't addressed directly, so just respond
|
|
topics = [x for x in what.split(' ') if len(x) >= 3]
|
|
|
|
return self.bot.reply(event, "{0:s}"
|
|
"".format(" ".join(markovlib.generate_line(context, topics=topics))))
|
|
|
|
def get_or_create_target_context(self, target_name):
|
|
"""Return the context for a provided nick/channel, creating missing ones."""
|
|
target_name = target_name.lower()
|
|
|
|
# find the stuff, or create it
|
|
channel, c = IrcChannel.objects.get_or_create(name=target_name, server=self.connection.server_config)
|
|
try:
|
|
target = MarkovTarget.objects.get(channel=channel)
|
|
except MarkovTarget.DoesNotExist:
|
|
# we need to create a context and a target, and we have to make the context first
|
|
# make a context --- lacking a good idea, just create one with this target name until configured otherwise
|
|
context, c = MarkovContext.objects.get_or_create(name=target_name)
|
|
target, c = MarkovTarget.objects.get_or_create(name=target_name, context=context, channel=channel)
|
|
|
|
return target.context
|
|
|
|
try:
|
|
return target.context
|
|
except MarkovContext.DoesNotExist:
|
|
# make a context --- lacking a good idea, just create one with this target name until configured otherwise
|
|
context, c = MarkovContext.objects.get_or_create(name=target_name)
|
|
target.context = context
|
|
target.save()
|
|
|
|
return target.context
|
|
|
|
|
|
plugin = Markov
|