lib/lita/handlers/markov/engine.rb in lita-markov-1.0.1 vs lib/lita/handlers/markov/engine.rb in lita-markov-1.0.2

@@ -67,12 +67,14 @@
           dictionary.insert entry.merge(frequency: 1)
         end
       end
     end
 
-    def random_capitalized_word
-      states = @db[:dictionary].map(:current_state)
+    def random_capitalized_word(user)
+      states = @db[:dictionary]
+        .where(user: user)
+        .map(:current_state)
 
       capitalized_states = states.select do |state|
         /^[A-Z]/ =~ state
       end
 
@@ -85,21 +87,22 @@
       raise EmptyDictionaryError, 'No data for user' if state.nil?
 
       return state.split(' ').first
     end
 
-    def random_second_word(first_word)
+    def random_second_word(user, first_word)
       states = @db[:dictionary]
         .where(Sequel.like(:current_state, first_word+'%'))
+        .where(user: user)
         .map(:current_state)
 
       state = states.sample
       state.split(' ').last
     end
 
     def is_punctuation?(string)
-      PUNCTUATION.any? { |p| string.end_with? p }
+      PUNCTUATION.any? { |p| string == p }
     end
 
     def get_next_state(user, current_state)
       states = @db[:dictionary]
         .where(user: user, current_state: current_state)
@@ -112,14 +115,15 @@
 
       distribution.sample
     end
 
     def generate_sentence_for(user, length = 30)
-      first_word = random_capitalized_word
-      second_word = random_second_word first_word
+      first_word = random_capitalized_word user
+      second_word = random_second_word user, first_word
 
       sentence = [first_word, second_word]
+      ended_with_punctuation = false
 
       while sentence.length < length
         current_state = sentence.slice(sentence.length - @depth, @depth).join ' '
 
         next_state = get_next_state user, current_state
@@ -127,13 +131,20 @@
         # Stop if we failed to find a next state
         break if next_state.nil?
 
         sentence << next_state
 
-        break if is_punctuation? next_state
+        if is_punctuation? next_state
+          ended_with_punctuation = true
+          break
+        end
       end
 
-      sentence.slice(0..-2).join(' ') + sentence.last
+      chain = sentence.slice(0..-2).join(' ')
+      chain << ' ' unless ended_with_punctuation
+      chain << sentence.last
+
+      chain
     end
 
     def separate_string string
       # Including the punctuation in group so they'll be included in the
       # split results