stroopwafel/Chatbot.py

import asyncio
import discord
from discord.ext import commands
import numpy as np

class Chatbot(object):
	"""Een heuse chatbot. Under construction..."""
	
	def __init__(self, bot):
		self.bot = bot
		self.voice_states = {}
		self.END_OF_MESSAGE = "<EOM>"
		try:
			np.load('markov_dict.npy')
		except FileNotFoundError:
			#If there's no initial file, make a new one:
			initFile = {'pairs':{'Hoi':{self.END_OF_MESSAGE : 1}}, 'trairs':{}}
			np.save("markov_dict.npy", initFile)
		except:
			print("ERROR! - Could not open or create Markov chat file")
	def __check(self, ctx):
		return True
		
	#TODO laat Stroop reageren op persoonlijke berichten
	
	#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, wordDict, message: str):
		newPairDict = wordDict['pairs']
		newTrairDict = wordDict['trairs']
		
		message = message + " " + self.END_OF_MESSAGE
		corpus = message.split()
		pairs = self.makePairs(corpus)
		trairs = self.makeTrairs(corpus)
		print(trairs)

		#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 {'pairs':newPairDict, 'trairs':newTrairDict};
				
	@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
		"""
		
		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:
		iterator = channel.history(limit=nrOfMessages, before=None, after=None, reverse=False, around=None)
		for i in range(nrOfMessages):
			msg = yield from iterator.next()
			messageContent = msg.content.strip()
			if messageContent == "":
				pass #print("-- Bericht is leeg. Negeer.")
			elif msg.author.bot:
				pass #print("-- Dit is een bot chat. Negeer.")
			elif messageContent[0] == "!":
				pass #print("-- Dit is een commando. Negeer.")
			elif "```" in messageContent:
				pass #print("-- Daar zit een codeblok in. Negeer.")
			else:
				wordDict = self.learnFrom(wordDict, messageContent)
			
		np.save("markov_dict.npy", wordDict)
		yield from ctx.channel.send("Ik heb weer wat nieuwe woordjes geleerd. :slight_smile:")
		
	
	@commands.command(pass_context=True, hidden=False)
	@asyncio.coroutine	
	def reageer(self, ctx):
		"""UNDER CONSTRUCTION. Laat Stroop iets willekeurigs zeggen.
		Stroop leert van onze gesprekken. Hij kan zich nu voordoen als een van ons. Soort van.
		"""
		
		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']
		#TODO store in global variable to minimise IO?

		#Delete command
		try:
			yield from ctx.message.delete()
		except:
			print("No permission to delete message")
			
		#print(wordDict)
		
		#Pick a random first word (Which is not an end-of-message)
		#TODO pick a word which responds to the previous message in the channel
		first_word = np.random.choice(list(wordPairDict.keys()))
		while first_word == self.END_OF_MESSAGE: #or first_word.islower() ???
			first_word = np.random.choice(list(wordPairDict.keys()))
		chain = [first_word]

		#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()
			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))

		chain = chain[:-1] #remove <EOM>
			
		#print(' '.join(chain))
		yield from ctx.channel.send(' '.join(chain))