From 8c5a446b754f0563b8c6af08f79204135be85c08 Mon Sep 17 00:00:00 2001 From: skiesel Date: Mon, 9 May 2016 16:45:02 -0400 Subject: [PATCH 1/2] markov babbler plugin --- main.go | 2 + plugins/babbler/babbler.go | 202 ++++++++++++++++++++++++++++++++ plugins/babbler/babbler_test.go | 73 ++++++++++++ 3 files changed, 277 insertions(+) create mode 100644 plugins/babbler/babbler.go create mode 100644 plugins/babbler/babbler_test.go diff --git a/main.go b/main.go index 238ff46..9b9b5e3 100644 --- a/main.go +++ b/main.go @@ -17,6 +17,7 @@ import ( "github.com/velour/catbase/plugins/fact" "github.com/velour/catbase/plugins/leftpad" "github.com/velour/catbase/plugins/reminder" + "github.com/velour/catbase/plugins/babbler" "github.com/velour/catbase/plugins/talker" "github.com/velour/catbase/plugins/your" "github.com/velour/catbase/slack" @@ -53,6 +54,7 @@ func main() { b.AddHandler("your", your.New(b)) b.AddHandler("counter", counter.New(b)) b.AddHandler("reminder", reminder.New(b)) + b.AddHandler("babbler", babbler.New(b)) // catches anything left, will always return true b.AddHandler("factoid", fact.New(b)) diff --git a/plugins/babbler/babbler.go b/plugins/babbler/babbler.go new file mode 100644 index 0000000..144b435 --- /dev/null +++ b/plugins/babbler/babbler.go @@ -0,0 +1,202 @@ +// © 2013 the CatBase Authors under the WTFPL. See AUTHORS for the list of authors. + +package babbler + +import ( + // "database/sql" + "fmt" + "math/rand" + "strings" + + + "github.com/jmoiron/sqlx" + "github.com/velour/catbase/bot" + "github.com/velour/catbase/bot/msg" +) + +type BabblerPlugin struct { + Bot bot.Bot + db *sqlx.DB + babblers map[string]*babbler +} + +type babbler struct { + start *node + end *node + lookup map[string]*node +} + +type node struct { + wordFrequency int + arcs map[string]*arc +} + +type arc struct { + transitionFrequency int + next *node +} + +func New(bot bot.Bot) *BabblerPlugin { + plugin := &BabblerPlugin{ + Bot: bot, + db: bot.DB(), + babblers: map[string]*babbler{}, + } + + // this who string isn't escaped, just sooo, you know. + babbler, err := getMarkovChain(plugin.db, "seabass") + if err == nil { + plugin.babblers["seabass"] = babbler + } else { + plugin.babblers["seabass"] = newBabbler() + } + + return plugin +} + +func (p *BabblerPlugin) Message(message msg.Message) bool { + lowercase := strings.ToLower(message.Body) + tokens := strings.Fields(lowercase) + + if _, ok := p.babblers[message.User.Name]; ok { + addToMarkovChain(p.babblers[message.User.Name], lowercase) + } + + if len(tokens) == 4 && strings.Contains(lowercase, "initialize babbler for ") { + who := tokens[len(tokens)-1] + if _, ok := p.babblers[who]; !ok { + babbler, err := getMarkovChain(p.db, who) + if err == nil { + p.babblers[who] = babbler + } else { + p.babblers[who] = newBabbler() + } + } + } + + if len(tokens) == 2 && tokens[1] == "says" { + if _, ok := p.babblers[tokens[0]]; ok { + p.Bot.SendMessage(message.Channel, p.babble(tokens[0])) + return true + } + } + return false +} + +func (p *BabblerPlugin) Help(channel string, parts []string) { + p.Bot.SendMessage(channel, "seabass says") +} + +func (p *BabblerPlugin) Event(kind string, message msg.Message) bool { + return false +} + +func (p *BabblerPlugin) BotMessage(message msg.Message) bool { + return false +} + +func (p *BabblerPlugin) RegisterWeb() *string { + return nil +} + +func addToMarkovChain(babble *babbler, phrase string) { + words := strings.Fields(strings.ToLower(phrase)) + + prev := babble.start + prev.wordFrequency++ + for i := range words { + // has this word been seen before + if _, ok := babble.lookup[words[i]]; !ok { + babble.lookup[words[i]] = &node{ + wordFrequency: 1, + arcs: map[string]*arc{}, + } + } else { + babble.lookup[words[i]].wordFrequency++ + } + + // has this word been seen after the previous word before + if _, ok := prev.arcs[words[i]]; !ok { + prev.arcs[words[i]] = &arc{ + transitionFrequency: 1, + next: babble.lookup[words[i]], + } + } else { + prev.arcs[words[i]].transitionFrequency++ + } + prev = babble.lookup[words[i]] + } + + // has this word ended a fact before + if _, ok := prev.arcs[""]; !ok { + prev.arcs[""] = &arc{ + transitionFrequency: 1, + next: babble.end, + } + } else { + prev.arcs[""].transitionFrequency++ + } +} + +func newBabbler() *babbler { + return &babbler { + start: &node { + wordFrequency: 0, + arcs: map[string]*arc{}, + }, + end: &node { + wordFrequency: 0, + arcs: map[string]*arc{}, + }, + lookup: map[string]*node{}, + } +} + +// this who string isn't escaped, just sooo, you know. +func getMarkovChain(db *sqlx.DB, who string) (*babbler, error) { + query := fmt.Sprintf(`select tidbit from factoid where fact like '%s quotes';`, who) + rows, err := db.Query(query) + if err != nil { + return nil, err + } + + babble := newBabbler() + + for rows.Next() { + + var tidbit string + err := rows.Scan(&tidbit) + if err != nil { + return nil, err + } + + addToMarkovChain(babble, tidbit) + } + return babble, nil +} + +func (p *BabblerPlugin) babble(who string) string { + if babbler, ok := p.babblers[who]; ok { + if len(babbler.start.arcs) == 0 { + return "" + } + words := []string{} + cur := babbler.start + for cur != babbler.end { + which := rand.Intn(cur.wordFrequency) + sum := 0 + for word, arc := range cur.arcs { + sum += arc.transitionFrequency + if sum > which { + words = append(words, word) + cur = arc.next + break + } + } + } + + return strings.Join(words, " ") + } + + return fmt.Sprintf("could not find babbler: %s", who) +} \ No newline at end of file diff --git a/plugins/babbler/babbler_test.go b/plugins/babbler/babbler_test.go new file mode 100644 index 0000000..da93362 --- /dev/null +++ b/plugins/babbler/babbler_test.go @@ -0,0 +1,73 @@ +// © 2013 the CatBase Authors under the WTFPL. See AUTHORS for the list of authors. + +package babbler + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/velour/catbase/bot" + "github.com/velour/catbase/bot/msg" + "github.com/velour/catbase/bot/user" +) + +func makeMessage(payload string) msg.Message { + isCmd := strings.HasPrefix(payload, "!") + if isCmd { + payload = payload[1:] + } + return msg.Message{ + User: &user.User{Name: "tester"}, + Channel: "test", + Body: payload, + Command: isCmd, + } +} + +func TestBabbler(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + assert.NotNil(t, c) + seabass := makeMessage("This is a message") + seabass.User = &user.User{Name: "seabass"} + res := c.Message(seabass) + seabass.Body = "This is another message" + res = c.Message(seabass) + seabass.Body = "This is a long message" + res = c.Message(seabass) + res = c.Message(makeMessage("!seabass says")) + assert.Len(t, mb.Messages, 1) + assert.True(t, res) + assert.Contains(t, mb.Messages[0], "this is") + assert.Contains(t, mb.Messages[0], "message") +} + +func TestHelp(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + assert.NotNil(t, c) + c.Help("channel", []string{}) + assert.Len(t, mb.Messages, 1) +} + +func TestBotMessage(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + assert.NotNil(t, c) + assert.False(t, c.BotMessage(makeMessage("test"))) +} + +func TestEvent(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + assert.NotNil(t, c) + assert.False(t, c.Event("dummy", makeMessage("test"))) +} + +func TestRegisterWeb(t *testing.T) { + mb := bot.NewMockBot() + c := New(mb) + assert.NotNil(t, c) + assert.Nil(t, c.RegisterWeb()) +} From e6324ad5a2b276e3dd4d9732303842aa40ee9c92 Mon Sep 17 00:00:00 2001 From: Chris Sexton Date: Tue, 10 May 2016 21:15:52 -0400 Subject: [PATCH 2/2] babbler: Add all users in chan, add cfg'd users Also misc small updates to the mock object for easier config usage --- bot/mock.go | 4 +- config/config.go | 3 + example_config.json | 5 ++ plugins/babbler/babbler.go | 98 ++++++++++++++++++++------------- plugins/babbler/babbler_test.go | 2 + slack/slack.go | 11 +++- 6 files changed, 82 insertions(+), 41 deletions(-) diff --git a/bot/mock.go b/bot/mock.go index 1655155..193c5ee 100644 --- a/bot/mock.go +++ b/bot/mock.go @@ -16,11 +16,13 @@ type MockBot struct { mock.Mock db *sqlx.DB + Cfg config.Config + Messages []string Actions []string } -func (mb *MockBot) Config() *config.Config { return &config.Config{} } +func (mb *MockBot) Config() *config.Config { return &mb.Cfg } func (mb *MockBot) DBVersion() int64 { return 1 } func (mb *MockBot) DB() *sqlx.DB { return mb.db } func (mb *MockBot) Who(string) []user.User { return []user.User{} } diff --git a/config/config.go b/config/config.go index df4ee2b..112d279 100644 --- a/config/config.go +++ b/config/config.go @@ -64,6 +64,9 @@ type Config struct { QuoteTime int StartupFact string } + Babbler struct { + DefaultUsers []string + } } // Readconfig loads the config data out of a JSON file located in cfile diff --git a/example_config.json b/example_config.json index 08eb01b..3540c3f 100644 --- a/example_config.json +++ b/example_config.json @@ -58,5 +58,10 @@ "QuoteTime": 1, "StartupFact": "speed test", "MinLen": 5 + }, + "Babbler": { + "DefaultUsers": [ + "seabass" + ] } } diff --git a/plugins/babbler/babbler.go b/plugins/babbler/babbler.go index 144b435..685bd04 100644 --- a/plugins/babbler/babbler.go +++ b/plugins/babbler/babbler.go @@ -8,59 +8,79 @@ import ( "math/rand" "strings" - "github.com/jmoiron/sqlx" "github.com/velour/catbase/bot" "github.com/velour/catbase/bot/msg" + "github.com/velour/catbase/bot/user" + "github.com/velour/catbase/config" ) type BabblerPlugin struct { - Bot bot.Bot - db *sqlx.DB + Bot bot.Bot + db *sqlx.DB + config *config.Config babblers map[string]*babbler } type babbler struct { - start *node - end *node + start *node + end *node lookup map[string]*node } type node struct { wordFrequency int - arcs map[string]*arc + arcs map[string]*arc } type arc struct { transitionFrequency int - next *node + next *node } func New(bot bot.Bot) *BabblerPlugin { plugin := &BabblerPlugin{ - Bot: bot, - db: bot.DB(), + Bot: bot, + db: bot.DB(), + config: bot.Config(), babblers: map[string]*babbler{}, } - // this who string isn't escaped, just sooo, you know. - babbler, err := getMarkovChain(plugin.db, "seabass") - if err == nil { - plugin.babblers["seabass"] = babbler - } else { - plugin.babblers["seabass"] = newBabbler() - } - return plugin } +func (p *BabblerPlugin) makeBabbler(newUser user.User) { + name := newUser.Name + babbler, err := getMarkovChain(p.db, name) + if err == nil { + p.babblers[name] = babbler + } else { + p.babblers[name] = newBabbler() + } +} + +func (p *BabblerPlugin) makeBabblers(newUser user.User) { + users := p.Bot.Who(p.config.MainChannel) + users = append(users, newUser) + for _, name := range p.config.Babbler.DefaultUsers { + users = append(users, user.New(name)) + } + for _, u := range users { + p.makeBabbler(u) + } +} + func (p *BabblerPlugin) Message(message msg.Message) bool { + if len(p.babblers) == 0 { + p.makeBabblers(*message.User) + } else if _, ok := p.babblers[message.User.Name]; !ok { + p.makeBabbler(*message.User) + } + lowercase := strings.ToLower(message.Body) tokens := strings.Fields(lowercase) - if _, ok := p.babblers[message.User.Name]; ok { - addToMarkovChain(p.babblers[message.User.Name], lowercase) - } + addToMarkovChain(p.babblers[message.User.Name], lowercase) if len(tokens) == 4 && strings.Contains(lowercase, "initialize babbler for ") { who := tokens[len(tokens)-1] @@ -71,20 +91,24 @@ func (p *BabblerPlugin) Message(message msg.Message) bool { } else { p.babblers[who] = newBabbler() } + p.Bot.SendMessage(message.Channel, "Okay.") + return true } } if len(tokens) == 2 && tokens[1] == "says" { - if _, ok := p.babblers[tokens[0]]; ok { - p.Bot.SendMessage(message.Channel, p.babble(tokens[0])) - return true + saying := p.babble(tokens[0]) + if saying == "" { + p.Bot.SendMessage(message.Channel, "Ze ain't said nothin'") } + p.Bot.SendMessage(message.Channel, saying) + return true } return false } func (p *BabblerPlugin) Help(channel string, parts []string) { - p.Bot.SendMessage(channel, "seabass says") + p.Bot.SendMessage(channel, "initialize babbler for seabass\n\nseabass says") } func (p *BabblerPlugin) Event(kind string, message msg.Message) bool { @@ -109,7 +133,7 @@ func addToMarkovChain(babble *babbler, phrase string) { if _, ok := babble.lookup[words[i]]; !ok { babble.lookup[words[i]] = &node{ wordFrequency: 1, - arcs: map[string]*arc{}, + arcs: map[string]*arc{}, } } else { babble.lookup[words[i]].wordFrequency++ @@ -119,7 +143,7 @@ func addToMarkovChain(babble *babbler, phrase string) { if _, ok := prev.arcs[words[i]]; !ok { prev.arcs[words[i]] = &arc{ transitionFrequency: 1, - next: babble.lookup[words[i]], + next: babble.lookup[words[i]], } } else { prev.arcs[words[i]].transitionFrequency++ @@ -131,7 +155,7 @@ func addToMarkovChain(babble *babbler, phrase string) { if _, ok := prev.arcs[""]; !ok { prev.arcs[""] = &arc{ transitionFrequency: 1, - next: babble.end, + next: babble.end, } } else { prev.arcs[""].transitionFrequency++ @@ -139,15 +163,15 @@ func addToMarkovChain(babble *babbler, phrase string) { } func newBabbler() *babbler { - return &babbler { - start: &node { - wordFrequency: 0, - arcs: map[string]*arc{}, - }, - end: &node { - wordFrequency: 0, - arcs: map[string]*arc{}, - }, + return &babbler{ + start: &node{ + wordFrequency: 0, + arcs: map[string]*arc{}, + }, + end: &node{ + wordFrequency: 0, + arcs: map[string]*arc{}, + }, lookup: map[string]*node{}, } } @@ -199,4 +223,4 @@ func (p *BabblerPlugin) babble(who string) string { } return fmt.Sprintf("could not find babbler: %s", who) -} \ No newline at end of file +} diff --git a/plugins/babbler/babbler_test.go b/plugins/babbler/babbler_test.go index da93362..aabbd7d 100644 --- a/plugins/babbler/babbler_test.go +++ b/plugins/babbler/babbler_test.go @@ -28,10 +28,12 @@ func makeMessage(payload string) msg.Message { func TestBabbler(t *testing.T) { mb := bot.NewMockBot() c := New(mb) + c.config.Babbler.DefaultUsers = []string{"seabass"} assert.NotNil(t, c) seabass := makeMessage("This is a message") seabass.User = &user.User{Name: "seabass"} res := c.Message(seabass) + assert.Len(t, c.babblers, 1) seabass.Body = "This is another message" res = c.Message(seabass) seabass.Body = "This is a long message" diff --git a/slack/slack.go b/slack/slack.go index 10adcc1..ead635f 100644 --- a/slack/slack.go +++ b/slack/slack.go @@ -258,9 +258,14 @@ func (s *Slack) Who(id string) []string { u := s.url + "channels.info" resp, err := http.PostForm(u, url.Values{"token": {s.config.Slack.Token}, "channel": {id}}) - if err != nil || resp.StatusCode != 200 { - log.Printf("Error posting user info request: %d %s", - resp.StatusCode, err) + if err != nil { + log.Printf("Error posting user info request: %s", + err) + return []string{} + } + if resp.StatusCode != 200 { + log.Printf("Error posting user info request: %d", + resp.StatusCode) return []string{} } defer resp.Body.Close()