stroopwafel/Chatbot.py

import asyncio
import discord
from discord.ext import commands
import numpy as np
import random
from datetime import datetime
import re

class Chatbot(object):
	"""Een heuse chatbot. Under construction..."""
	
	def __init__(self, bot):
		self.bot = bot
		self.voice_states = {}
		self.END_OF_MESSAGE = "<EOM>"
		self.START_OF_MESSAGE = "<SOM>"
		self.pmLast = datetime(2018, 1, 1, 1, 1, 1)
		self.pmCount = 0
		try:
			np.load('markov_dict.npy')
		except FileNotFoundError:
			#If there's no initial file, make a new one:
			initFile = {'pairs':{self.START_OF_MESSAGE:{'Hoi':1},'Hoi':{self.END_OF_MESSAGE:1}}, 'trairs':{}, 'responsePairs':{}}
			np.save("markov_dict.npy", initFile)
		except:
			print("ERROR! - Could not open or create Markov chat file")
			
			
		#Do stuff when a message comes in:
		@self.bot.event
		@asyncio.coroutine
		def on_message(message):
			if message.author.bot: 
				return #ignore your own messages, Stroop
			
			if message.content.strip() == "" or message.content.strip()[0] != '!': #don't respond to a command
				if message.content.lower() == "nihao":
					yield from message.channel.send("kankerlauw")
					return
				if message.content.lower() == "lol":
                                        yield from message.channel.send("laffing aut laut")
                                        return
				if "DMChannel" in type(message.channel).__name__: #respond to messages in PM channels
					if (self.isAllowedToLearnFromPm(message)): #avoid spam-learning
						self.pmCount += 1
						self.pmLast = datetime.now()
						yield from self.learnFromMessagesInChannel(message.channel.id, 1)
					yield from self.respondToMessage(message)
				elif message.channel.id == self.bot.hdChannels["belangrijk"]: #don't abuse the "belangrijk" channel
					pass
				else:
					a = random.randint(0,999)
					if message.channel.id == self.bot.hdChannels["shitpost"] and a < 40:  #shitpost channel
						yield from self.learnFromMessagesInChannel(message.channel.id, 10)
						yield from self.respondToMessage(message)
					elif a < 10:                                             #all other channels
						yield from self.learnFromMessagesInChannel(message.channel.id, 20)
						yield from self.respondToMessage(message)

			#process commands:
			yield from self.bot.process_commands(message)
			
	def __check(self, ctx):
		return True
		
	
	#The "Guus filter"
	#Don't allow to learn from very large personal messages
	#Only learn from 10 personal messages per hour at most
	def isAllowedToLearnFromPm(self, message):
		if len(message.content.strip()) > 150:
			return False
		diff = datetime.now() - self.pmLast
		if diff.days > 0 or diff.seconds > 3600:
			self.pmCount = 0
			return True
		if self.pmCount > 10:
			return False
		return True
	
	#Make word pairs:
	def makePairs(self, corpus):
		for i in range(len(corpus)-1):
			yield (corpus[i], corpus[i+1])
		
	#Make word trairs. Yes, I completely made up "trair"
	#It's like pairs, but with three, but not really, because there are two. OK?
	#[1 0 1]
	def makeTrairs(self, corpus):
		for i in range(1, len(corpus)-1):
			yield (corpus[i-1], corpus[i+1])
		
	#Learn from a human message
	#Returns an updated dict
	def learnFrom(self, wordDictPairs, wordDictTrairs, message: str):
		newPairDict = wordDictPairs
		newTrairDict = wordDictTrairs
		
		#If the message is only a tag, don't learn from it. return
		if message.count(' ') == 0 and message[0:2] == "<@":
			return newPairDict, newTrairDict;
		
		message = self.START_OF_MESSAGE + " " + message + " " + self.END_OF_MESSAGE
		corpus = message.split()
		for n, word in enumerate(corpus):
			if word == "@everyone":
				corpus[n] = "@iedereen"
		pairs = self.makePairs(corpus)
		trairs = self.makeTrairs(corpus)

		#Add wordcounts to pair dict
		for word_1, word_2 in pairs:
			if word_1 in newPairDict.keys():
				if word_2 in newPairDict[word_1].keys():
					newPairDict[word_1][word_2] += 1
				else:
					newPairDict[word_1][word_2] = 1
			else:
				newPairDict[word_1] = {word_2 : 1}		
		#Add wordcounts to trair dict
		for word_1, word_3 in trairs:
			if word_1 in newTrairDict.keys():
				if word_3 in newTrairDict[word_1].keys():
					newTrairDict[word_1][word_3] += 1
				else:
					newTrairDict[word_1][word_3] = 1
			else:
				newTrairDict[word_1] = {word_3 : 1}
				
		return newPairDict, newTrairDict;
	
	#Learn to respond to messages
	#Returns a new dict with response pairs
	def learnResponse(self, wordDictResp, message1: str, message2: str):
		STOP_WORDS = ["aan", "achter", "al", "ben", "dan", "dat", "de", "die", "dit", "een", "en", "er", "gehad", "had", "heb", 
			"hen", "het", "hun", "in", "is", "kon", "met", "na", "of", "om", "ook", "op", "tot", "van", "was", "wat", "zo", 
			"a", "an", "the", "am", "was", "were", "be", "not", "at"]
			
		newRespPairDict = wordDictResp
		
		#Split messages into words and remove stopwords:
		message1List = []
		message2List = []
		for word in message1.split():
			if word not in STOP_WORDS:
				message1List.append(word)
		for word in message2.split():
			if word not in STOP_WORDS:
				message2List.append(word)
		
		#link words of response to words of initial message
		for word_1 in message1List:
			for word_2 in message2List:
				if word_1 in newRespPairDict.keys():
					if word_2 in newRespPairDict[word_1].keys():
						newRespPairDict[word_1][word_2] += 1
					else:
						newRespPairDict[word_1][word_2] = 1
				else:
					newRespPairDict[word_1] = {word_2 : 1}	

		return newRespPairDict

			

	@commands.command(pass_context=True, hidden=True)
	@asyncio.coroutine	
	def trainmarkovchat(self, ctx, channelId=409751773595566081, nrOfMessages=50):
		"""Train that chain, jermaine
		param channelId: use this channel to train
		param nrOfMessages: use the x most recent messages
		"""
		
		yield from self.learnFromMessagesInChannel(channelId, nrOfMessages)
		
		yield from ctx.channel.send("Ik heb weer wat nieuwe woordjes geleerd. :slight_smile:")

	
	#Learns from an x number of messages from a certain channel
	@asyncio.coroutine
	def learnFromMessagesInChannel(self, channelId, nrOfMessages):
		
		channel=self.bot.get_channel(channelId)
		#Open current chain:
		wordDict = np.load('markov_dict.npy').item()
		
		#Iterate through the last x messages from the channel:
		messageIterator = channel.history(limit=nrOfMessages, before=None, after=None, reverse=False, around=None)
		allMessages = []
		for i in range(nrOfMessages):
			msg = yield from messageIterator.next()
			allMessages.append(msg)
		
		i = 0
		nextMessageIndex = -10
		#Note: the loop goes from new to old
		for msg in allMessages:
			messageContent = msg.content.replace('(', '').replace(')', '').replace('"', '')
			messageContent = messageContent.strip()

			if messageContent == "":
				pass # Bericht is leeg. Negeer.
			elif msg.author.bot:
				pass # Dit is een bot chat. Negeer.
			elif messageContent[0] == "!":
				pass # Dit is een commando. Negeer.
			elif "```" in messageContent:
				pass # Daar zit een codeblok in. Negeer.
			elif messageContent[0:7] == "http://" or messageContent[0:8] == "https://":
				pass # Een link. Negeer. TODO: post af en toe een random link die geen repost is
			else:
				# Usable message!
				wordDict['pairs'], wordDict['trairs'] = self.learnFrom(wordDict['pairs'], wordDict['trairs'], messageContent)
				
				if (i-nextMessageIndex) == 1 and msg.author != nextMessage.author: #the previous usable messages was the one right after this one and was not from the same author
					timeDelta = nextMessage.created_at - msg.created_at
					#TODO check for different author
					if timeDelta.days == 0 and timeDelta.seconds < 3600: #less than an hour difference
						#this is counted as a response. learn from it
						wordDict['responsePairs'] = self.learnResponse(wordDict['responsePairs'], messageContent, nextMessageContent)
				
				#save for next iteration:
				nextMessage = msg
				nextMessageContent = messageContent
				nextMessageIndex = i
			i += 1
			
		np.save("markov_dict.npy", wordDict)
		
	
	@commands.command(pass_context=True, hidden=False)
	@asyncio.coroutine	
	def reageer(self, ctx):
		"""Laat Stroop ergens op reageren.
		Stroop leert van onze gesprekken. Hij kan zich nu voordoen als een van ons. Soort van.
		"""
		
		#Get previous message to respond to
		messageIterator = ctx.channel.history(limit=3, before=None, after=None, reverse=False, around=None)
		messageToRespondTo = yield from messageIterator.next() #this will be the !reageer command. skip
		messageToRespondTo = yield from messageIterator.next()
		if messageToRespondTo.content.strip()[0] == '!' or messageToRespondTo.author.bot: #another chance
			messageToRespondTo = yield from messageIterator.next() #but it's not a huge problem if he responds to himself
		if messageToRespondTo.content.strip()[0] == '!': 
			wordsToRespondTo = []
		else:
			wordsToRespondTo = messageToRespondTo.content.strip().split()
		#print("respond to: ", wordsToRespondTo)

		#Delete command
		try:
			yield from ctx.message.delete()
		except:
			pass
			
		response = self.getResponseForMessage(wordsToRespondTo)
		yield from ctx.channel.send(response)
	
	#Respond to a message.
	#Assumes the message isn't a command
	@asyncio.coroutine
	def respondToMessage(self, message):
		wordsToRespondTo = message.content.strip().split()
		
		response = self.getResponseForMessage(wordsToRespondTo)
		yield from message.channel.send(response)
		
	#Gets a response
	def getResponseForMessage(self, wordsToRespondTo):
		
		MAX_N_WORDS = 100 #Don't use more words than this
		
		#Get dictionary/Markov chain
		wordDict = np.load('markov_dict.npy').item()
		wordPairDict = wordDict['pairs']
		wordTrairDict = wordDict['trairs']
		responsePairDict = wordDict['responsePairs']
		#TODO store in global variable to minimise IO?
		
		#Start with <SOM>
		chain = [self.START_OF_MESSAGE]

		#Go through the chain:
		for i in range(1, MAX_N_WORDS):
			
			if chain[-1] == self.END_OF_MESSAGE:
				break
			
			#Find the possibilities:
			pairPossibilities = wordPairDict[chain[-1]].copy()
			
			#Increase score if it would be a good response:
			for userWord in wordsToRespondTo:
				if userWord in responsePairDict:
					for goodResponseWord in responsePairDict[userWord]:
						if goodResponseWord in pairPossibilities.keys():
							if i == 1:
								pairPossibilities[goodResponseWord] += responsePairDict[userWord][goodResponseWord] * 2 #first word? try to respond
							else:
								pairPossibilities[goodResponseWord] += responsePairDict[userWord][goodResponseWord] * 0.1
			#Increase score for trairs:
			if i > 1: #we can only compare trairs if we have at least two words already
				if chain[-2] in wordTrairDict:
					trairPossibilities = wordTrairDict[chain[-2]].copy()
					#If a trair matches, increase the frequency of the pair
					for word2 in pairPossibilities.keys():
						if word2 in trairPossibilities.keys():
							pairPossibilities[word2] += trairPossibilities[word2]
							
			#Calculate probabilities and pick a next word:
			#https://stackoverflow.com/questions/835092/python-dictionary-are-keys-and-values-always-the-same-order
			possible_words = list(pairPossibilities.keys())
			weights = np.array(list(pairPossibilities.values()))
			probs = weights / weights.sum()
			chain.append(np.random.choice(possible_words, p=probs))
	

		if chain[0] == self.START_OF_MESSAGE:
			chain = chain[1:-1] #remove <SOM> and <EOM>
		else:
			chain = chain[:-1] #remove <EOM>
		
		#remove 90% of tags:
		for n, word in enumerate(chain):
			if word[0:2] == "<@":
				a = random.randint(1,10)
				if a != 1:
					chain[n] = "@jemoeder"
			
		#print(' '.join(chain))
		#yield from ctx.channel.send(' '.join(chain))
		return ' '.join(chain)
		
		
		
		
	#Try to count the number of syllables in a word
	#Used to make Haikus
	def countSyllables(self, word: str):
		KNOWN_ACRONYMS = ["aub", "svp", "ivm", "nr", "zsm", "mbt", "tmi", "fyi", "afk", "omg", "idk", "coc"]
		word = word.lower().strip()
		
		#for acryonyms, just count the letters:
		if word in KNOWN_ACRONYMS or re.match("^([a-z]{1}[.]{1}){2,}$", word) or re.match("^([bcdfghjklmnpqrstvwxz]){2,}$", word):
			return len(word.replace(".", ""))
		
		#count vowels surrounded by consonants
		count = len(re.findall("[bcdfghjklmnpqrstvwxyz]{1}[aeiouyëäïöüéèáà]{1,2}[bcdfghjklmnpqrstvwxz]{0,1}|[aeiouyëäïöüéèáà]{1,2}[bcdfghjklmnpqrstvwxz]{1}", word))
		count += len(re.findall("[ëäï]", word)) #ideeën etc
		
		return max(1, count)
		
	#Append a message change with a numer of syllables
	#If it returns False, it couldn't find anything
	#It's recursive, yo
	def appendMessageWithNrOfSyllables(self, nrOfSyl: int, chain, wordPairDict):	
		#Find the possibilities:
		if chain[-1] == "\n":
			pairPossibilities = wordPairDict[chain[-2]].copy()
		else:
			pairPossibilities = wordPairDict[chain[-1]].copy()
		
		for poss in list(pairPossibilities.keys()):
			if self.countSyllables(poss) > nrOfSyl or poss == self.END_OF_MESSAGE:  #TODO change to constant
				del pairPossibilities[poss]
				
		if len(pairPossibilities) == 0:
			return False
		
		#Calculate probabilities and pick a next word:
		possible_words = list(pairPossibilities.keys())
		weights = np.array(list(pairPossibilities.values()))
		probs = weights / weights.sum()
		chain.append(np.random.choice(possible_words, p=probs))
		#TODO this is duplicate code. reuse the stuff from getResponseForMessage()
		
		if self.countSyllables(chain[-1]) == nrOfSyl:
			return chain
		#RECURSION!
		return self.appendMessageWithNrOfSyllables( (nrOfSyl-self.countSyllables(chain[-1])) , chain, wordPairDict)
		
	#Make a nice Haiku	
	def getHaiku(self):
		MAX_N_TRIES = 50
		
		#Get dictionary/Markov chain
		wordDict = np.load('markov_dict.npy').item()
		wordPairDict = wordDict['pairs']

		#try:
		for i in range(1, MAX_N_TRIES):
			chain = [self.START_OF_MESSAGE]
			
			chain = self.appendMessageWithNrOfSyllables(5, chain, wordPairDict)
			if chain is not False:
				chain.append("\n")
				chain = self.appendMessageWithNrOfSyllables(7, chain, wordPairDict)
			if chain is not False:
				chain.append("\n")
				chain = self.appendMessageWithNrOfSyllables(5, chain, wordPairDict)
			if chain is not False:
				break
				
		chain = chain[1:] #remove <SOM>
		return ' '.join(chain)
		
	@commands.command(pass_context=True, hidden=False)
	@asyncio.coroutine	
	def haiku(self, ctx):
		"""Stroop is een poëet. Laat hem een Haiku maken.
		"""
			
		response = self.getHaiku()
		yield from ctx.channel.send(response)