chatterbot 0.2.0

data/lib/chatterbot/db.rb

@@ -0,0 +1,70 @@
+module Chatterbot
+  
+  #
+  # routines for optionally interacting with a database for logging
+  # tweets, and storing config data there. Uses Sequel to handle the
+  # heavy lifing.
+  module DB   
+    #
+    # connect to the database, and generate any missing tables
+    def db
+      @_db ||= connect_and_validate
+    end
+
+    protected  
+
+    #
+    # get a DB object from Sequel
+    def get_connection
+      if has_sequel?
+        Sequel.connect(config[:db_uri])
+      end
+    end
+    
+    #
+    # try and connect to the DB, and create tables that are missing.
+    def connect_and_validate
+      conn = get_connection
+      return if conn.nil?
+
+      if ! conn.tables.include?(:blacklist)
+        conn.create_table :blacklist do
+          String :user, :primary_key => true
+          DateTime :created_at
+        end
+      end
+      
+      if ! conn.tables.include?(:tweets)
+        conn.create_table :tweets do
+          primary_key :id
+          String :txt
+          String :bot
+          String :user
+          String :source_id
+          String :source_tweet
+          
+          DateTime :created_at
+        end
+      end
+
+      if ! conn.tables.include?(:config)
+        conn.create_table :config do
+          String :id, :primary_key => true
+
+          Bignum :since_id
+
+          String :secret
+          String :token
+          String :consumer_secret
+          String :consumer_key
+          
+          DateTime :created_at
+          DateTime :updated_at
+        end
+      end
+      
+      conn
+    end
+
+  end
+end

data/lib/chatterbot/dsl.rb

@@ -0,0 +1,100 @@
+require File.join(File.dirname(__FILE__), '..', 'chatterbot')
+require 'optparse'
+
+module Chatterbot
+  #
+  # very basic DSL to handle the common stuff you would want to do with a bot.
+  module DSL
+
+    #
+    # generate a Bot object. if the DSL is being called from a Bot object, just return it
+    # otherwise create a bot and return that
+    def bot
+      return @bot unless @bot.nil?
+      
+      #
+      # parse any command-line options and use them to initialize the bot
+      #
+      params = {}
+      
+      opts = OptionParser.new
+
+      opts.banner = "Usage: #{File.basename($0)} [options]"
+
+      opts.separator ""
+      opts.separator "Specific options:" 
+      
+      opts.on('-d', '--db [ARG]', "Specify a DB connection URI")    { |d| ENV["chatterbot_db"] = d }
+      opts.on('-c', '--config [ARG]', "Specify a config file to use")    { |c| ENV["chatterbot_config"] = c }
+      opts.on('-t', '--test', "Run the bot without actually sending any tweets") { params[:debug_mode] = true }
+      opts.on('--dry-run', "Run the bot in test mode, and also don't update the database")    { params[:debug_mode] = true ; params[:no_update] = true }
+      opts.on('-s', '--since_id [ARG]', "Check for tweets since tweet id #[ARG]")    { |s| params[:since_id] = s }
+
+      opts.on_tail("-h", "--help", "Show this message") do
+        puts opts
+        exit
+      end      
+      
+      opts.parse!(ARGV)
+
+      @bot = Chatterbot::Bot.new(params)
+    end
+
+    #
+    # specify a bot-specific blacklist of users.  accepts an array, or a 
+    # comma-delimited string
+    def blacklist(b=nil)
+      if b.is_a?(String)
+        b = b.split(",").collect { |s| s.strip }
+      end
+
+      if b.nil? || b.empty?
+        bot.blacklist = []
+      else     
+        bot.blacklist += b
+      end
+    
+    end
+    
+    #
+    # specify list of strings we will check when deciding to respond to a tweet or not
+    def exclude(e=nil)
+      if e.is_a?(String)
+        e = e.split(",").collect { |s| s.strip }
+      end
+
+      if e.nil? || e.empty?
+        bot.exclude = []
+      else     
+        bot.exclude += e
+      end
+    end
+    
+    #
+    # search twitter for the specified terms
+    def search(query, &block)
+      @bot.search(query, &block)
+    end
+    
+    #
+    # handle replies to the bot
+    def replies(&block)
+      @bot.replies(&block)
+    end
+
+    #
+    # send a tweet
+    def tweet(txt, params = {}, original = nil)
+      @bot.tweet(txt, params, original)
+    end   
+
+    #
+    # reply to a tweet
+    def reply(txt, source)
+      @bot.reply(txt, source)
+    end
+  end
+end
+
+include Chatterbot::DSL
+include Chatterbot::Helpers

data/lib/chatterbot/helpers.rb

@@ -0,0 +1,29 @@
+module Chatterbot
+
+  #
+  # a bunch of helper routines for bots
+  module Helpers
+
+    #
+    # The name of the currently running bot
+    def botname
+      if self.class < Bot
+        self.class.to_s.downcase
+      else
+        File.basename($0,".rb")
+      end
+    end
+
+    #
+    # Take the incoming tweet/user name, and turn it into something suitable for replying 
+    # to a user. Basically, get their handle and add a '@' to it.
+    def tweet_user(tweet)     
+      if ! tweet.is_a?(String)
+        base = tweet.has_key?(:from_user) ? tweet[:from_user] : tweet[:user][:screen_name]
+      else
+        base = tweet
+      end
+      base =~ /^@/ ? base : "@#{base}"
+    end
+  end
+end

data/lib/chatterbot/logging.rb

@@ -0,0 +1,48 @@
+require 'logger'
+
+module Chatterbot
+
+  #
+  # routines for outputting log messages, as well as logging tweets
+  # to the database if desired.
+  module Logging
+
+    #
+    # log a message
+    def debug(s)
+      logger.debug s unless ! logging?
+    end
+
+    #
+    # something really bad happened, print it out and log it
+    def critical(s)
+      puts s
+      debug s
+    end
+    
+    #
+    # log a tweet to the database
+    def log(txt, source=nil)
+      return unless log_tweets?
+
+      data = {:txt => txt, :bot => botname, :created_at => Time.now}
+      if source != nil
+        data = data.merge(:user => source[:from_user],
+                          :source_id => source[:id],
+                          :source_tweet => source[:text])
+      end
+
+      # populate the table
+      db[:tweets].insert(data)
+    end
+
+protected
+    #
+    # initialize a Logger object, writing to log_dest
+    def logger
+      # log to the dest specified in the config file, rollover after 10mb of data
+      @_logger ||= Logger.new(log_dest, 0, 1024 * 1024)
+    end
+
+  end
+end

data/lib/chatterbot/reply.rb

@@ -0,0 +1,28 @@
+module Chatterbot
+
+  #
+  # reply method for responding to tweets
+  module Reply
+
+    # handle replies for the bot
+    def replies(&block)
+      return unless require_login
+      debug "check for replies since #{since_id}"
+      
+      results = client.replies(:since_id => since_id)
+
+      if results.is_a?(Hash) && results.has_key?("error")
+        critical results["error"]
+      else
+        results.each { |s|
+          s.symbolize_keys!
+          unless ! block_given? || on_blacklist?(s) || skip_me?(s)
+            update_since_id(s)
+            yield s         
+          end
+        }
+      end
+    end
+
+  end
+end

data/lib/chatterbot/search.rb

@@ -0,0 +1,35 @@
+module Chatterbot
+
+  #
+  # handle Twitter searches
+  module Search
+
+    # internal search code
+    def search(queries, &block)
+      return unless init_client
+      
+      debug "check for tweets since #{since_id}"
+
+      if queries.is_a?(String)
+        queries = [queries]
+      end
+      
+      #
+      # search twitter
+      #
+      queries.each { |query|
+        search = client.search(query, default_opts)
+        update_since_id(search)
+
+        if search != nil
+          search["results"].each { |s|
+            s.symbolize_keys!
+            yield s unless ! block_given? || on_blacklist?(s) || skip_me?(s)
+          }
+        end
+      }
+    end
+  
+  end
+end
+

data/lib/chatterbot/tweet.rb

@@ -0,0 +1,25 @@
+module Chatterbot
+
+  #
+  # routines for sending tweets
+  module Tweet 
+   
+    # simple wrapper for sending a message
+    def tweet(txt, params = {}, original = nil)
+      return if require_login == false
+
+      if debug_mode?
+        debug "I'm in debug mode, otherwise I would tweet: #{txt}"
+      else
+        debug txt
+        log txt, original
+        client.update txt, params
+      end
+    end   
+
+    # reply to a tweet
+    def reply(txt, source)
+      self.tweet txt, {:in_reply_to_status_id => source["id"]}, source
+    end   
+  end
+end

data/lib/chatterbot/version.rb

@@ -0,0 +1,3 @@
+module Chatterbot
+  VERSION = "0.2.0"
+end

data/spec/blacklist_spec.rb

@@ -0,0 +1,115 @@
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
+
+describe "Chatterbot::Blacklist" do
+  it "has a list of excluded phrases" do
+    @bot = test_bot
+    @bot.exclude = ["junk", "i hate bots", "foobar", "spam"]
+    @bot.skip_me?("did you know that i hate bots?").should == true
+  end
+
+  describe "skip_me?" do
+    before(:each) do
+      @bot = test_bot
+      @bot.stub!(:exclude).and_return(["junk", "i hate bots", "foobar", "spam"])
+    end
+
+    it "blocks tweets with phrases we don't want" do
+      @bot.skip_me?("did you know that i hate bots?").should == true
+    end
+
+    it "isn't case-specific" do
+      @bot.skip_me?("DID YOU KNOW THAT I HATE BOTS?").should == true
+    end
+
+
+    it "allows users we do want" do
+      @bot.skip_me?("a tweet without any bad content").should == false
+    end
+
+    it "works with result hashes" do
+      @bot.skip_me?({:text => "did you know that i hate bots?"}).should == true
+      @bot.skip_me?({:text => "a tweet without any bad content"}).should == false
+    end
+  end
+
+  describe "on_blacklist?" do
+    before(:each) do
+      @bot = test_bot
+      @bot.stub!(:blacklist).and_return(["skippy", "blippy", "dippy"])
+    end
+
+    it "blocks users we don't want" do
+      @bot.on_blacklist?("skippy").should == true
+    end
+
+    it "allows users we do want" do
+      @bot.on_blacklist?("flippy").should == false
+    end
+
+    it "isn't case-specific" do
+      @bot.on_blacklist?("FLIPPY").should == false
+      @bot.on_blacklist?("SKIPPY").should == true
+    end
+
+    it "works with result hashes" do
+      @bot.on_blacklist?({:from_user => "skippy"}).should == true
+      @bot.on_blacklist?({:from_user => "flippy"}).should == false
+    end   
+  end
+
+  describe "on_global_blacklist?" do
+    before(:each) do
+      @bot = test_bot
+    end
+
+    it "doesn't freak out if no db" do
+      @bot.should_receive(:has_db?).and_return(false)
+      @bot.on_global_blacklist?("foobar").should == false
+    end
+
+    it "collects name from the db if it exists" do
+      @bot.stub!(:has_db?).and_return(true)
+      blacklist_table = mock(Object)
+      mock_dataset = mock(Object, {:count => 1})
+      blacklist_table.should_receive(:where).
+        with({ :user => "a"}).
+        and_return( mock_dataset )
+
+      
+      missing_dataset = mock(Object, {:count => 0})
+      blacklist_table.should_receive(:where).
+        with({ :user => "b"}).
+        and_return( missing_dataset )
+      
+      @bot.stub!(:db).and_return({ 
+                                   :blacklist => blacklist_table
+                                 })
+      @bot.on_global_blacklist?("a").should == true
+      @bot.on_global_blacklist?("b").should == false
+    end
+  end
+
+
+  describe "db interaction" do
+    before(:each) do
+      @db_tmp = Tempfile.new("blacklist.db")
+      @db_uri = "sqlite://#{@db_tmp.path}"
+
+      @bot = Chatterbot::Bot.new    
+      @bot.config[:db_uri] = @db_uri
+      @bot.db
+    end
+
+    describe "add_to_blacklist" do
+      it "adds to the blacklist table" do      
+        @bot.add_to_blacklist("tester")
+      end
+
+      it "doesn't add a double entry" do
+        @bot.add_to_blacklist("tester")
+        @bot.add_to_blacklist("tester")      
+      end
+    end
+    
+  end
+end