Brian Stephan
333424025b
markov targets are queried and autogenerated based on chatter, but had a legacy name which is no longer in use for this, preferring the foreign keys to channel and consequently server. the name is really just informative these days, but was still being used to find targets, and thus was breaking when two servers had the same channel name in them. this fixes that
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
|