Remove twitter storage adapter in favor of trainer

chatterbot/adapters/storage/__init__.py

@@ -2,4 +2,3 @@
 from .django_storage import DjangoStorageAdapter
 from .jsonfile import JsonFileStorageAdapter
 from .mongodb import MongoDatabaseAdapter
-from .twitter_storage import TwitterAdapter

chatterbot/chatterbot.py

@@ -70,7 +70,7 @@ def __init__(self, name, **kwargs):
         # Use specified trainer or fall back to the default
         trainer = kwargs.get('trainer', 'chatterbot.trainers.Trainer')
         TrainerClass = import_module(trainer)
-        self.trainer = TrainerClass(self.storage)
+        self.trainer = TrainerClass(self.storage, **kwargs)
 
         self.logger = kwargs.get('logger', logging.getLogger(__name__))
 

chatterbot/trainers.py

@@ -1,12 +1,14 @@
 from .conversation import Statement, Response
 from .corpus import Corpus
+import logging
 
 
 class Trainer(object):
 
     def __init__(self, storage, **kwargs):
         self.storage = storage
         self.corpus = Corpus()
+        self.logger = logging.getLogger(__name__)
 
     def train(self, *args, **kwargs):
         raise self.TrainerInitializationException()
@@ -82,3 +84,89 @@ def train(self, *corpora):
             for data in corpus_data:
                 for pair in data:
                     trainer.train(pair)
+
+
+class TwitterTrainer(Trainer):
+
+    def __init__(self, storage, **kwargs):
+        super(TwitterTrainer, self).__init__(storage, **kwargs)
+        from twitter import Api as TwitterApi
+
+        self.api = TwitterApi(
+            consumer_key=kwargs.get('twitter_consumer_key'),
+            consumer_secret=kwargs.get('twitter_consumer_secret'),
+            access_token_key=kwargs.get('twitter_access_token_key'),
+            access_token_secret=kwargs.get('twitter_access_token_secret')
+        )
+
+    def random_word(self, base_word='random'):
+        """
+        Generate a random word using the Twitter API.
+
+        Search twitter for recent tweets containing the term 'random'.
+        Then randomly select one word from those tweets and do another
+        search with that word. Return a randomly selected word from the
+        new set of results.
+        """
+        import random
+        random_tweets = self.api.GetSearch(term=base_word, count=5)
+        random_words = self.get_words_from_tweets(random_tweets)
+        random_word = random.choice(list(random_words))
+        tweets = self.api.GetSearch(term=random_word, count=5)
+        words = self.get_words_from_tweets(tweets)
+        word = random.choice(list(words))
+        return word
+
+    def get_words_from_tweets(self, tweets):
+        """
+        Given a list of tweets, return the set of
+        words from the tweets.
+        """
+        words = set()
+
+        for tweet in tweets:
+            # TODO: Handle non-ascii characters properly
+            cleaned_text = ''.join(
+                [i if ord(i) < 128 else ' ' for i in tweet.text]
+            )
+            tweet_words = cleaned_text.split()
+
+            for word in tweet_words:
+                # If the word contains only letters with a length from 4 to 9
+                if word.isalpha() and len(word) > 3 and len(word) <= 9:
+                    words.add(word)
+
+        return words
+
+    def get_statements(self):
+        """
+        Returns list of random statements from the API.
+        """
+        from twitter import TwitterError
+        statements = []
+
+        # Generate a random word
+        random_word = self.random_word()
+
+        self.logger.info(u'Requesting 50 random tweets containing the word {}'.format(random_word))
+        tweets = self.api.GetSearch(term=random_word, count=50)
+        for tweet in tweets:
+            statement = Statement(tweet.text)
+
+            if tweet.in_reply_to_status_id:
+                try:
+                    status = self.api.GetStatus(tweet.in_reply_to_status_id)
+                    statement.add_response(Response(status.text))
+                    statements.append(statement)
+                except TwitterError as e:
+                    self.logger.warning(str(e))
+
+        self.logger.info('Adding {} tweets with responses'.format(len(statements)))
+
+        return statements
+
+    def train(self):
+        for i in range(0, 10):
+            statements = self.get_statements()
+            for statement in statements:
+                self.storage.update(statement, force=True)

examples/twitter_example.py → examples/twitter_training_example.py

@@ -1,8 +1,12 @@
 from chatterbot import ChatBot
 from settings import TWITTER
+import logging
 
 
 '''
+This example demonstrates how you can train your chat bot
+using data from Twitter.
+
 To use this example, create a new file called settings.py.
 In settings.py define the following:
 
@@ -14,25 +18,23 @@
 }
 '''
 
-chatbot = ChatBot("ChatterBot",
-    storage_adapter="chatterbot.adapters.storage.TwitterAdapter",
+# Comment out the following line to disable verbose logging
+logging.basicConfig(level=logging.INFO)
+
+chatbot = ChatBot("TwitterBot",
     logic_adapters=[
         "chatterbot.adapters.logic.ClosestMatchAdapter"
     ],
     input_adapter="chatterbot.adapters.input.TerminalAdapter",
     output_adapter="chatterbot.adapters.output.TerminalAdapter",
-    database="../database.db",
+    database="./twitter-database.db",
     twitter_consumer_key=TWITTER["CONSUMER_KEY"],
     twitter_consumer_secret=TWITTER["CONSUMER_SECRET"],
     twitter_access_token_key=TWITTER["ACCESS_TOKEN"],
-    twitter_access_token_secret=TWITTER["ACCESS_TOKEN_SECRET"]
+    twitter_access_token_secret=TWITTER["ACCESS_TOKEN_SECRET"],
+    trainer="chatterbot.trainers.TwitterTrainer"
 )
 
-print("Type something to begin...")
-
-while True:
-    try:
-        bot_input = chatbot.get_response(None)
+chatbot.train()
 
-    except (KeyboardInterrupt, EOFError, SystemExit):
-        break
+chatbot.logger.info('Trained database generated successfully!')