gpt2: refactor

This commit is contained in:
Chris Sexton 2021-02-01 07:38:12 -05:00 committed by Chris Sexton
parent dfbcfafab9
commit b7576bd855
1 changed files with 10 additions and 14 deletions

View File

@ -5,6 +5,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"net/http" "net/http"
"regexp"
"strings" "strings"
"github.com/velour/catbase/bot" "github.com/velour/catbase/bot"
@ -23,27 +24,22 @@ func New(b bot.Bot) *GPT2Plugin {
c: b.Config(), c: b.Config(),
} }
b.Register(p, bot.Message, p.message) b.RegisterRegexCmd(p, bot.Message, gpt2Regex, p.gpt2Cmd)
b.Register(p, bot.Help, p.help) b.Register(p, bot.Help, p.help)
return p return p
} }
const prefix = "gpt2" var gpt2Regex = regexp.MustCompile(`(?i)^gpt2 (?P<input>.*)$`)
func (p *GPT2Plugin) message(c bot.Connector, kind bot.Kind, message msg.Message, args ...interface{}) bool { func (p *GPT2Plugin) gpt2Cmd(r bot.Request) bool {
ch := message.Channel input := r.Values["input"]
lowerBody := strings.ToLower(message.Body)
if message.Command && strings.HasPrefix(lowerBody, prefix) {
input := message.Body[len(prefix)+1:]
txt, err := p.getGPTText(input) txt, err := p.getGPTText(input)
if err != nil { if err != nil {
txt = p.c.Get("gpt.error", "The GPT service is unavailable.") txt = p.c.Get("gpt.error", "The GPT service is unavailable.")
} }
p.b.Send(c, bot.Message, ch, txt) p.b.Send(r.Conn, bot.Message, r.Msg.Channel, txt)
return true return true
}
return false
} }
func (p *GPT2Plugin) help(c bot.Connector, kind bot.Kind, message msg.Message, args ...interface{}) bool { func (p *GPT2Plugin) help(c bot.Connector, kind bot.Kind, message msg.Message, args ...interface{}) bool {