chatbot.py

from autobahn.twisted.websocket import WebSocketServerProtocol, WebSocketServerFactory, listenWS
from twisted.python import log
from autobahn.util import newid, utcnow
import Cookie
from chatterbot import ChatBot
from chatterbot.trainers import ChatterBotCorpusTrainer
from chatterbot.conversation import Statement
from users import Users

import sys
from twisted.python import log
from twisted.internet import reactor, defer
from twisted.internet.task import deferLater
from pymongo import MongoClient

"""
Database
"""
db_name = "chatty-database"
db_uri = "mongodb://localhost:27017/"
client = MongoClient(db_uri)
db = client[db_name]
# Users collection
users = db.users

"""
Create users class
"""
userList = Users(users)

"""
Create chatbot
"""
chatbot = ChatBot('Chatty',
  storage_adapter="chatterbot.adapters.storage.MongoDatabaseAdapter",
  logic_adapters=[
    "chatterbot.adapters.logic.ClosestMatchAdapter",
    "name_logic_adapter.NameLogicAdapter",
    "chatterbot.adapters.logic.MathematicalEvaluation",
    "chatterbot.adapters.logic.TimeLogicAdapter",
  ],
  # input_adapter="chatterbot.adapters.input.TerminalAdapter",
  output_adapter="output_adapter.OutputFormatAdapter",
  tie_breaking_method="random_response",
  database=db_name,
  trainer='chatterbot.trainers.ChatterBotCorpusTrainer',
  userList=userList,
  database_uri=db_uri,
  read_only=True,
  tokenizer=tokenize,
)

"""
Train the chatbot
"""
db['statements'].delete_many({})
chatbot.train(
  "corpus.english"
)

"""
Websocket
"""

class MyServerProtocol(WebSocketServerProtocol):
  """
  Set the user_id here
  """
  def onConnect(self, request):
    self._user_id = None
    self.userObj = {}
    if 'cookie' in request.headers:
      try:
        cookie = Cookie.SimpleCookie()
        cookie.load(str(request.headers['cookie']))
      except Cookie.CookieError:
        pass
      else:
        if 'user_id' in cookie:
          user_id = cookie['user_id'].value
          self._user_id = user_id
          log.msg("Cookie already set: %s" % self._user_id)

    if self._user_id is None:
      self._user_id = newid()
      maxAge = 86400
      user_id = {'created': utcnow(), 'authenticated': None, 'maxAge': maxAge, 'connections': set()}
      headers['Set-Cookie'] = 'user_id=%s;max-age=%d' % (self._user_id, maxAge)

    return (None, headers)

  def onOpen(self):
    # Add the user to the collection
    userFromDB = users.find_one({ "id": self._user_id })

    if userFromDB:
      self.userObj = { x: userFromDB[x] if userFromDB[x] != 'None' else None for x in userFromDB }
      self.userObj['isNew'] = False
      # Reset current intent
      self.userObj['current_intent'] = None
    else:
      self.userObj = {
        "id": self._user_id,
        "name": None,
        "email": None,
        "phone": None,
        "current_intent": None,
        "previous_intent": None,
        "isNew": True
      }
      users.insert_one(self.userObj)

    # Add user to userList object
    userList.add_user(self._user_id, self.userObj)
    """
    Send an initial message
    """
    name = self.userObj['name']
    if self.userObj['isNew'] != False or name == None:
      self.sendConversations([
        "Hello there",
        "I am Chatty, your friendly chatbot."
      ])
    else:
      self.sendConversations([
        "Welcome back, " + name,
        "Let me know how can i be of help to you!"
      ])

  @defer.inlineCallbacks
  def sendConversations(self, conversations):
    for conv in conversations:
      yield deferLater(reactor, 0.5, lambda: None)
      self.sendMessage(conv.encode('utf8'))

  @defer.inlineCallbacks
  def sendSingleMessage(self, msg):
    yield deferLater(reactor, 0.25, lambda: None)
    self.sendMessage(msg)

  def onMessage(self, payload, isBinary):
    """
    Add user_id to input statement
    """
    input = Statement(payload)
    input.add_extra_data('user', self.userObj)
    msg, confidence = chatbot.get_response(input)

    # Check for confidence here
    # If confidence score is very low send a `I dont understand response`
    if confidence < 0.3:
      return self.sendConversations([
        "I am not sure if I understand you properly.",
        "How can i help you?"
      ])

    # Check for multiple conversations
    conversations = msg.extra_data["conversations"] if "conversations" in msg.extra_data else None

    if conversations:
      self.sendConversations(conversations)
    else:
      self.sendSingleMessage(msg.text.encode('utf8'))

  def onClose(self, wasClean, code, reason):
    print("WebSocket connection closed: {0}".format(reason))

"""
Entry
"""
if __name__ == '__main__':
  log.startLogging(sys.stdout)
  headers = {}
  factory = WebSocketServerFactory(u"ws://127.0.0.1:9000", headers=headers)
  factory.protocol = MyServerProtocol
  listenWS(factory)
  reactor.run()

users.py

class Users():
  def __init__(self, collection):
    self._users = dict()
    self.collection = collection

  def add_user(self, userId, obj):
    self._users[userId] = obj

  def set_intent(self, userId, name):
    self._users[userId]['current_intent'] = name
    self.collection.find_one_and_update({ 'id': userId }, { '$set' : self._users[userId] })

  def update(self, userId, key, value):
    self._users[userId][key] = value
    self.collection.find_one_and_update({ 'id': userId }, { '$set' : self._users[userId] })

  def get_intent(self, userId):
    return self._users[userId]['current_intent']

  def remove_intent(self, userId):
    temp = self._users[userId]['current_intent']
    self._users[userId]['current_intent'] = None
    self._users[userId]['previous_intent'] = temp

    self.collection.find_one_and_update({ 'id': userId }, { '$set' : self._users[userId]})