batch learning for the babbler

plugins/babbler/babbler.go

@@ -3,7 +3,6 @@
 package babbler
 
 import (
-	// "database/sql"
 	"fmt"
 	"math/rand"
 	"strings"
@@ -87,8 +86,8 @@ func (p *BabblerPlugin) Message(message msg.Message) bool {
 		}
 		p.Bot.SendMessage(message.Channel, saying)
 		return true
-	} else if len(tokens) == 4 && strings.Contains(lowercase, "initialize babbler for ") {
+	} else if len(tokens) == 4 && strings.Index(lowercase, "initialize babbler for ") == 0 {
-		who := tokens[len(tokens)-1]
+		who := tokens[3]
 		if _, ok := p.babblers[who]; !ok {
 			babbler, err := getMarkovChain(p.db, who)
 			if err == nil {
@@ -99,6 +98,30 @@ func (p *BabblerPlugin) Message(message msg.Message) bool {
 			p.Bot.SendMessage(message.Channel, "Okay.")
 			return true
 		}
+	} else if strings.Index(lowercase, "batch learn for ") == 0 {
+		who := tokens[3]
+		if _, ok := p.babblers[who]; !ok {
+			p.babblers[who] = newBabbler()
+		}
+
+		body := strings.Join(tokens[4:], " ")
+		body = strings.ToLower(body)
+
+		for _, a := range strings.Split(body, ".") {
+			for _, b := range strings.Split(a, "!") {
+				for _, c := range strings.Split(b, "?") {
+					for _, d := range strings.Split(c, "\n") {
+						trimmed := strings.TrimSpace(d)
+						if trimmed != "" {
+							addToMarkovChain(p.babblers[who], trimmed)
+						}
+					}
+				}
+			}
+		}
+
+		p.Bot.SendMessage(message.Channel, "Phew that was tiring.")
+		return true
 	} else {
 		addToMarkovChain(p.babblers[message.User.Name], lowercase)
 	}

plugins/babbler/babbler_test.go

@@ -45,6 +45,22 @@ func TestBabbler(t *testing.T) {
 	assert.Contains(t, mb.Messages[0], "message")
 }
 
+func TestBabblerBatch(t *testing.T) {
+	mb := bot.NewMockBot()
+	c := New(mb)
+	c.config.Babbler.DefaultUsers = []string{"seabass"}
+	assert.NotNil(t, c)
+	seabass := makeMessage("batch learn for seabass This is a message! This is another message. This is not a long message? This is not a message! This is not another message. This is a long message?")
+	res := c.Message(seabass)
+	assert.Len(t, c.babblers, 2)
+	assert.Len(t, mb.Messages, 1)
+	res = c.Message(makeMessage("!seabass says"))
+	assert.Len(t, mb.Messages, 2)
+	assert.True(t, res)
+	assert.Contains(t, mb.Messages[1], "this is")
+	assert.Contains(t, mb.Messages[1], "message")
+}
+
 func TestHelp(t *testing.T) {
 	mb := bot.NewMockBot()
 	c := New(mb)