2016-05-09 20:45:02 +00:00
|
|
|
// © 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"
|
2016-05-11 01:15:52 +00:00
|
|
|
"github.com/velour/catbase/bot/user"
|
|
|
|
"github.com/velour/catbase/config"
|
2016-05-09 20:45:02 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
type BabblerPlugin struct {
|
2016-05-11 01:15:52 +00:00
|
|
|
Bot bot.Bot
|
|
|
|
db *sqlx.DB
|
|
|
|
config *config.Config
|
2016-05-09 20:45:02 +00:00
|
|
|
babblers map[string]*babbler
|
|
|
|
}
|
|
|
|
|
|
|
|
type babbler struct {
|
2016-05-11 01:15:52 +00:00
|
|
|
start *node
|
|
|
|
end *node
|
2016-05-09 20:45:02 +00:00
|
|
|
lookup map[string]*node
|
|
|
|
}
|
|
|
|
|
|
|
|
type node struct {
|
|
|
|
wordFrequency int
|
2016-05-11 01:15:52 +00:00
|
|
|
arcs map[string]*arc
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
type arc struct {
|
|
|
|
transitionFrequency int
|
2016-05-11 01:15:52 +00:00
|
|
|
next *node
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func New(bot bot.Bot) *BabblerPlugin {
|
|
|
|
plugin := &BabblerPlugin{
|
2016-05-11 01:15:52 +00:00
|
|
|
Bot: bot,
|
|
|
|
db: bot.DB(),
|
|
|
|
config: bot.Config(),
|
2016-05-09 20:45:02 +00:00
|
|
|
babblers: map[string]*babbler{},
|
|
|
|
}
|
|
|
|
|
2016-05-11 01:15:52 +00:00
|
|
|
return plugin
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *BabblerPlugin) makeBabbler(newUser user.User) {
|
|
|
|
name := newUser.Name
|
|
|
|
babbler, err := getMarkovChain(p.db, name)
|
2016-05-09 20:45:02 +00:00
|
|
|
if err == nil {
|
2016-05-11 01:15:52 +00:00
|
|
|
p.babblers[name] = babbler
|
2016-05-09 20:45:02 +00:00
|
|
|
} else {
|
2016-05-11 01:15:52 +00:00
|
|
|
p.babblers[name] = newBabbler()
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
2016-05-11 01:15:52 +00:00
|
|
|
}
|
2016-05-09 20:45:02 +00:00
|
|
|
|
2016-05-11 01:15:52 +00:00
|
|
|
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)
|
|
|
|
}
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (p *BabblerPlugin) Message(message msg.Message) bool {
|
2016-05-11 01:15:52 +00:00
|
|
|
if len(p.babblers) == 0 {
|
|
|
|
p.makeBabblers(*message.User)
|
|
|
|
} else if _, ok := p.babblers[message.User.Name]; !ok {
|
|
|
|
p.makeBabbler(*message.User)
|
|
|
|
}
|
|
|
|
|
2016-05-09 20:45:02 +00:00
|
|
|
lowercase := strings.ToLower(message.Body)
|
|
|
|
tokens := strings.Fields(lowercase)
|
|
|
|
|
2016-05-11 01:15:52 +00:00
|
|
|
addToMarkovChain(p.babblers[message.User.Name], lowercase)
|
2016-05-09 20:45:02 +00:00
|
|
|
|
|
|
|
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()
|
|
|
|
}
|
2016-05-11 01:15:52 +00:00
|
|
|
p.Bot.SendMessage(message.Channel, "Okay.")
|
|
|
|
return true
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if len(tokens) == 2 && tokens[1] == "says" {
|
2016-05-11 01:15:52 +00:00
|
|
|
saying := p.babble(tokens[0])
|
|
|
|
if saying == "" {
|
|
|
|
p.Bot.SendMessage(message.Channel, "Ze ain't said nothin'")
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
2016-05-11 01:15:52 +00:00
|
|
|
p.Bot.SendMessage(message.Channel, saying)
|
|
|
|
return true
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
return false
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *BabblerPlugin) Help(channel string, parts []string) {
|
2016-05-11 01:15:52 +00:00
|
|
|
p.Bot.SendMessage(channel, "initialize babbler for seabass\n\nseabass says")
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
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,
|
2016-05-11 01:15:52 +00:00
|
|
|
arcs: map[string]*arc{},
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
} 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,
|
2016-05-11 01:15:52 +00:00
|
|
|
next: babble.lookup[words[i]],
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
} 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,
|
2016-05-11 01:15:52 +00:00
|
|
|
next: babble.end,
|
2016-05-09 20:45:02 +00:00
|
|
|
}
|
|
|
|
} else {
|
|
|
|
prev.arcs[""].transitionFrequency++
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func newBabbler() *babbler {
|
2016-05-11 01:15:52 +00:00
|
|
|
return &babbler{
|
|
|
|
start: &node{
|
|
|
|
wordFrequency: 0,
|
|
|
|
arcs: map[string]*arc{},
|
|
|
|
},
|
|
|
|
end: &node{
|
|
|
|
wordFrequency: 0,
|
|
|
|
arcs: map[string]*arc{},
|
|
|
|
},
|
2016-05-09 20:45:02 +00:00
|
|
|
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)
|
2016-05-11 01:15:52 +00:00
|
|
|
}
|