data/rbot/plugins/markov.rb

   1 #-- vim:sw=2:et
   2 #++
   3 #
   4 # :title: Markov plugin
   5 #
   6 # Author:: Tom Gilbert <tom@linuxbrit.co.uk>
   7 # Copyright:: (C) 2005 Tom Gilbert
   8 #
   9 # Contribute to chat with random phrases built from word sequences learned
  10 # by listening to chat
  11
  12 class MarkovPlugin < Plugin
  13   Config.register Config::BooleanValue.new('markov.enabled',
  14     :default => false,
  15     :desc => "Enable and disable the plugin")
  16   Config.register Config::IntegerValue.new('markov.probability',
  17     :default => 25,
  18     :validate => Proc.new { |v| (0..100).include? v },
  19     :desc => "Percentage chance of markov plugin chipping in")
  20   Config.register Config::ArrayValue.new('markov.ignore',
  21     :default => [],
  22     :desc => "Hostmasks and channel names markov should NOT learn from (e.g. idiot*!*@*, #privchan).")
  23   Config.register Config::IntegerValue.new('markov.max_words',
  24     :default => 50,
  25     :validate => Proc.new { |v| (0..100).include? v },
  26     :desc => "Maximum number of words the bot should put in a sentence")
  27   Config.register Config::IntegerValue.new('markov.learn_delay',
  28     :default => 0.5,
  29     :validate => Proc.new { |v| v >= 0 },
  30     :desc => "Time the learning thread spends sleeping after learning a line. If set to zero, learning from files can be very CPU intensive, but also faster.")
  31
  32   def initialize
  33     super
  34     @registry.set_default([])
  35     if @registry.has_key?('enabled')
  36       @bot.config['markov.enabled'] = @registry['enabled']
  37       @registry.delete('enabled')
  38     end
  39     if @registry.has_key?('probability')
  40       @bot.config['markov.probability'] = @registry['probability']
  41       @registry.delete('probability')
  42     end
  43     if @bot.config['markov.ignore_users']
  44       debug "moving markov.ignore_users to markov.ignore"
  45       @bot.config['markov.ignore'] = @bot.config['markov.ignore_users'].dup
  46       @bot.config.delete('markov.ignore_users'.to_sym)
  47     end
  48     @learning_queue = Queue.new
  49     @learning_thread = Thread.new do
  50       while s = @learning_queue.pop
  51         learn_line s
  52         sleep @bot.config['markov.learn_delay'] unless @bot.config['markov.learn_delay'].zero?
  53       end
  54     end
  55     @learning_thread.priority = -1
  56   end
  57
  58   def cleanup
  59     debug 'closing learning thread'
  60     @learning_queue.push nil
  61     @learning_thread.join
  62     debug 'learning thread closed'
  63   end
  64
  65   # if passed a pair, pick a word from the registry using the pair as key.
  66   # otherwise, pick a word from an given list
  67   def pick_word(word1, word2=:nonword)
  68     if word1.kind_of? Array
  69       wordlist = word1
  70     else
  71       wordlist = @registry["#{word1} #{word2}"]
  72     end
  73     wordlist.pick_one || :nonword
  74   end
  75
  76   def generate_string(word1, word2)
  77     # limit to max of markov.max_words words
  78     if word2
  79       output = "#{word1} #{word2}"
  80     else
  81       output = word1.to_s
  82     end
  83
  84     if @registry.key? output
  85       wordlist = @registry[output]
  86       wordlist.delete(:nonword)
  87     else
  88       output.downcase!
  89       keys = []
  90       @registry.each_key(output) do |key|
  91         if key.downcase.include? output
  92           keys << key
  93         else
  94           break
  95         end
  96       end
  97       if keys.empty?
  98         keys = @registry.keys.select { |k| k.downcase.include? output }
  99       end
 100       return nil if keys.empty?
 101       while key = keys.delete_one
 102         wordlist = @registry[key]
 103         wordlist.delete(:nonword)
 104         unless wordlist.empty?
 105           output = key
 106           word1, word2 = output.split
 107           break
 108         end
 109       end
 110     end
 111
 112     word3 = pick_word(wordlist)
 113     return nil if word3 == :nonword
 114
 115     output << " #{word3}"
 116     word1, word2 = word2, word3
 117
 118     (@bot.config['markov.max_words'] - 1).times do
 119       word3 = pick_word(word1, word2)
 120       break if word3 == :nonword
 121       output << " #{word3}"
 122       word1, word2 = word2, word3
 123     end
 124     return output
 125   end
 126
 127   def help(plugin, topic="")
 128     topic, subtopic = topic.split
 129
 130     case topic
 131     when "ignore"
 132       case subtopic
 133       when "add"
 134         "markov ignore add <hostmask|channel> => ignore a hostmask or a channel"
 135       when "list"
 136         "markov ignore list => show ignored hostmasks and channels"
 137       when "remove"
 138         "markov ignore remove <hostmask|channel> => unignore a hostmask or channel"
 139       else
 140         "ignore hostmasks or channels -- topics: add, remove, list"
 141       end
 142     when "status"
 143       "markov status => show if markov is enabled, probability and amount of messages in queue for learning"
 144     when "probability"
 145       "markov probability [<percent>] => set the % chance of rbot responding to input, or display the current probability"
 146     when "chat"
 147       case subtopic
 148       when "about"
 149         "markov chat about <word> [<another word>] => talk about <word> or riff on a word pair (if possible)"
 150       else
 151         "markov chat => try to say something intelligent"
 152       end
 153     else
 154       "markov plugin: listens to chat to build a markov chain, with which it can (perhaps) attempt to (inanely) contribute to 'discussion'. Sort of.. Will get a *lot* better after listening to a lot of chat. Usage: 'chat' to attempt to say something relevant to the last line of chat, if it can -- help topics: ignore, status, probability, chat, chat about"
 155     end
 156   end
 157
 158   def clean_str(s)
 159     str = s.dup
 160     str.gsub!(/^\S+[:,;]/, "")
 161     str.gsub!(/\s{2,}/, ' ') # fix for two or more spaces
 162     return str.strip
 163   end
 164
 165   def probability?
 166     return @bot.config['markov.probability']
 167   end
 168
 169   def status(m,params)
 170     if @bot.config['markov.enabled']
 171       reply = _("markov is currently enabled, %{p}% chance of chipping in") % { :p => probability? }
 172       l = @learning_queue.length
 173       reply << (_(", %{l} messages in queue") % {:l => l}) if l > 0
 174     else
 175       reply = _("markov is currently disabled")
 176     end
 177     m.reply reply
 178   end
 179
 180   def ignore?(m=nil)
 181     return false unless m
 182     return true if m.address? or m.private?
 183     @bot.config['markov.ignore'].each do |mask|
 184       return true if m.channel.downcase == mask.downcase
 185       return true if m.source.matches?(mask)
 186     end
 187     return false
 188   end
 189
 190   def ignore(m, params)
 191     action = params[:action]
 192     user = params[:option]
 193     case action
 194     when 'remove':
 195       if @bot.config['markov.ignore'].include? user
 196         s = @bot.config['markov.ignore']
 197         s.delete user
 198         @bot.config['ignore'] = s
 199         m.reply _("%{u} removed") % { :u => user }
 200       else
 201         m.reply _("not found in list")
 202       end
 203     when 'add':
 204       if user
 205         if @bot.config['markov.ignore'].include?(user)
 206           m.reply _("%{u} already in list") % { :u => user }
 207         else
 208           @bot.config['markov.ignore'] = @bot.config['markov.ignore'].push user
 209           m.reply _("%{u} added to markov ignore list") % { :u => user }
 210         end
 211       else
 212         m.reply _("give the name of a person or channel to ignore")
 213       end
 214     when 'list':
 215       m.reply _("I'm ignoring %{ignored}") % { :ignored => @bot.config['markov.ignore'].join(", ") }
 216     else
 217       m.reply _("have markov ignore the input from a hostmask or a channel. usage: markov ignore add <mask or channel>; markov ignore remove <mask or channel>; markov ignore list")
 218     end
 219   end
 220
 221   def enable(m, params)
 222     @bot.config['markov.enabled'] = true
 223     m.okay
 224   end
 225
 226   def probability(m, params)
 227     if params[:probability]
 228       @bot.config['markov.probability'] = params[:probability].to_i
 229       m.okay
 230     else
 231       m.reply _("markov has a %{prob}% chance of chipping in") % { :prob => probability? }
 232     end
 233   end
 234
 235   def disable(m, params)
 236     @bot.config['markov.enabled'] = false
 237     m.okay
 238   end
 239
 240   def should_talk
 241     return false unless @bot.config['markov.enabled']
 242     prob = probability?
 243     return true if prob > rand(100)
 244     return false
 245   end
 246
 247   def delay
 248     1 + rand(5)
 249   end
 250
 251   def random_markov(m, message)
 252     return unless should_talk
 253
 254     word1, word2 = message.split(/\s+/)
 255     return unless word1 and word2
 256     line = generate_string(word1, word2)
 257     return unless line
 258     # we do nothing if the line we return is just an initial substring
 259     # of the line we received
 260     return if message.index(line) == 0
 261     @bot.timer.add_once(delay) {
 262       m.reply line, :nick => false, :to => :public
 263     }
 264   end
 265
 266   def chat(m, params)
 267     line = generate_string(params[:seed1], params[:seed2])
 268     if line and line != [params[:seed1], params[:seed2]].compact.join(" ")
 269       m.reply line
 270     else
 271       m.reply _("I can't :(")
 272     end
 273   end
 274
 275   def rand_chat(m, params)
 276     # pick a random pair from the db and go from there
 277     word1, word2 = :nonword, :nonword
 278     output = Array.new
 279     @bot.config['markov.max_words'].times do
 280       word3 = pick_word(word1, word2)
 281       break if word3 == :nonword
 282       output << word3
 283       word1, word2 = word2, word3
 284     end
 285     if output.length > 1
 286       m.reply output.join(" ")
 287     else
 288       m.reply _("I can't :(")
 289     end
 290   end
 291
 292   def learn(*lines)
 293     lines.each { |l| @learning_queue.push l }
 294   end
 295
 296   def unreplied(m)
 297     return if ignore? m
 298
 299     # in channel message, the kind we are interested in
 300     message = clean_str m.plainmessage
 301
 302     if m.action?
 303       message = "#{m.sourcenick} #{message}"
 304     end
 305
 306     learn message
 307     random_markov(m, message) unless m.replied?
 308   end
 309
 310   def learn_triplet(word1, word2, word3)
 311       k = "#{word1} #{word2}"
 312       @registry[k] = @registry[k].push(word3)
 313   end
 314
 315   def learn_line(message)
 316     # debug "learning #{message}"
 317     wordlist = message.split(/\s+/)
 318     return unless wordlist.length >= 2
 319     word1, word2 = :nonword, :nonword
 320     wordlist << :nonword
 321     wordlist.each do |word3|
 322       learn_triplet(word1, word2, word3)
 323       word1, word2 = word2, word3
 324     end
 325   end
 326
 327   # TODO allow learning from URLs
 328   def learn_from(m, params)
 329     begin
 330       path = params[:file]
 331       file = File.open(path, "r")
 332       pattern = params[:pattern].empty? ? nil : Regexp.new(params[:pattern].to_s)
 333     rescue Errno::ENOENT
 334       m.reply _("no such file")
 335       return
 336     end
 337
 338     if file.eof?
 339       m.reply _("the file is empty!")
 340       return
 341     end
 342
 343     if params[:testing]
 344       lines = []
 345       range = case params[:lines]
 346       when /^\d+\.\.\d+$/
 347         Range.new(*params[:lines].split("..").map { |e| e.to_i })
 348       when /^\d+$/
 349         Range.new(1, params[:lines].to_i)
 350       else
 351         Range.new(1, [@bot.config['send.max_lines'], 3].max)
 352       end
 353
 354       file.each do |line|
 355         next unless file.lineno >= range.begin
 356         lines << line.chomp
 357         break if file.lineno == range.end
 358       end
 359
 360       lines = lines.map do |l|
 361         pattern ? l.scan(pattern).to_s : l
 362       end.reject { |e| e.empty? }
 363
 364       if pattern
 365         unless lines.empty?
 366           m.reply _("example matches for that pattern at lines %{range} include: %{lines}") % {
 367             :lines => lines.map { |e| Underline+e+Underline }.join(", "),
 368             :range => range.to_s
 369           }
 370         else
 371           m.reply _("the pattern doesn't match anything at lines %{range}") % {
 372             :range => range.to_s
 373           }
 374         end
 375       else
 376         m.reply _("learning from the file without a pattern would learn, for example: ")
 377         lines.each { |l| m.reply l }
 378       end
 379
 380       return
 381     end
 382
 383     if pattern
 384       file.each { |l| learn(l.scan(pattern).to_s) }
 385     else
 386       file.each { |l| learn(l.chomp) }
 387     end
 388
 389     m.okay
 390   end
 391 end
 392
 393 plugin = MarkovPlugin.new
 394 plugin.map 'markov ignore :action :option', :action => "ignore"
 395 plugin.map 'markov ignore :action', :action => "ignore"
 396 plugin.map 'markov ignore', :action => "ignore"
 397 plugin.map 'markov enable', :action => "enable"
 398 plugin.map 'markov disable', :action => "disable"
 399 plugin.map 'markov status', :action => "status"
 400 plugin.map 'chat about :seed1 [:seed2]', :action => "chat"
 401 plugin.map 'chat', :action => "rand_chat"
 402 plugin.map 'markov probability [:probability]', :action => "probability",
 403            :requirements => {:probability => /^\d+%?$/}
 404 plugin.map 'markov learn from :file [:testing [:lines lines]] [using pattern *pattern]', :action => "learn_from", :thread => true,
 405            :requirements => {
 406              :testing => /^testing$/,
 407              :lines   => /^(?:\d+\.\.\d+|\d+)$/ }
 408
 409 plugin.default_auth('ignore', false)
 410 plugin.default_auth('probability', false)
 411 plugin.default_auth('learn', false)
 412