// © 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" "github.com/velour/catbase/bot/user" "github.com/velour/catbase/config" ) type BabblerPlugin struct { Bot bot.Bot db *sqlx.DB config *config.Config 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(), config: bot.Config(), babblers: map[string]*babbler{}, } 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) 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() } p.Bot.SendMessage(message.Channel, "Okay.") return true } } if len(tokens) == 2 && tokens[1] == "says" { 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, "initialize babbler for seabass\n\nseabass 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) }