120 lines
3.2 KiB
Go
120 lines
3.2 KiB
Go
package main
|
|
|
|
import (
|
|
"Kevin/pkg/llama"
|
|
"Kevin/pkg/structs"
|
|
"fmt"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/bwmarrin/discordgo"
|
|
"github.com/spf13/viper"
|
|
)
|
|
|
|
var (
|
|
messages []structs.Message = []structs.Message{}
|
|
brain *llama.LLama
|
|
config *viper.Viper
|
|
)
|
|
|
|
func main() {
|
|
config = viper.New()
|
|
config.SetConfigType("toml")
|
|
config.SetConfigName("kevin")
|
|
config.AddConfigPath(".")
|
|
config.SetDefault("prefix", "!")
|
|
config.SetDefault("model", "models/ggml-vicuna-7b-1.1-q4_0.bin")
|
|
config.SetDefault("name", "Kevin")
|
|
config.SetDefault("discord.token", "not-set")
|
|
config.SetDefault("discord.chat", "not-set")
|
|
config.SetDefault("discord.owner", "not-set")
|
|
err := config.ReadInConfig()
|
|
if err != nil {
|
|
config.WriteConfigAs("kevin.toml")
|
|
panic(fmt.Errorf("fatal error config file: %s", err))
|
|
}
|
|
|
|
if config.GetString("discord.token") == "not-set" {
|
|
panic(fmt.Errorf("fatal error config file: %s", "discord token not set"))
|
|
}
|
|
|
|
brain, err = llama.New(config.GetString("model"), llama.SetContext(128), llama.SetParts(-1))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
discord, err := discordgo.New("Bot " + config.GetString("discord.token"))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
discord.AddHandler(messageCreate)
|
|
discord.AddHandler(ready)
|
|
discord.Identify.Intents = discordgo.MakeIntent(discordgo.IntentsGuildMessages)
|
|
|
|
err = discord.Open()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
fmt.Println("Bot is now running. Press CTRL-C to exit.")
|
|
sc := make(chan os.Signal, 1)
|
|
signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
|
|
<-sc
|
|
discord.Close()
|
|
brain.Free()
|
|
}
|
|
|
|
func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
|
|
if m.ChannelID != config.GetString("discord.chat") {
|
|
return
|
|
}
|
|
if m.Author.ID == s.State.User.ID {
|
|
return
|
|
}
|
|
if strings.HasPrefix(m.Content, config.GetString("prefix")) {
|
|
if m.Author.ID == config.GetString("discord.owner") {
|
|
if strings.HasPrefix(m.Content, "!clear") {
|
|
messages = []structs.Message{}
|
|
s.ChannelMessageSend(config.GetString("discord.chat"), "> Cleared messages")
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
// logic
|
|
s.UpdateStatusComplex(generateStatus(string(discordgo.StatusDoNotDisturb)))
|
|
message := structs.NewMessage(m.Author.Username, strings.TrimSpace(m.Content))
|
|
prompt := structs.GeneratePrompt(config, messages, *message)
|
|
var seed int = int(time.Now().Unix())
|
|
res, err := brain.Predict(prompt, llama.SetThreads(4), llama.SetTokens(128), llama.SetTopK(90), llama.SetTopP(0.86), llama.SetSeed(seed), llama.SetStopPrompt("###"))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
if res == "" {
|
|
res = "> Something went wrong"
|
|
}
|
|
s.ChannelMessageSend(config.GetString("discord.chat"), res)
|
|
message.Response = strings.TrimSpace(res)
|
|
messages = append(messages, *message)
|
|
s.UpdateStatusComplex(generateStatus(string(discordgo.StatusOnline)))
|
|
}
|
|
|
|
func ready(s *discordgo.Session, r *discordgo.Ready) {
|
|
s.UpdateStatusComplex(generateStatus(string(discordgo.StatusOnline)))
|
|
s.ChannelMessageSend(config.GetString("discord.chat"), "> Kevin is now running")
|
|
}
|
|
|
|
func generateStatus(status string) discordgo.UpdateStatusData {
|
|
return discordgo.UpdateStatusData{
|
|
AFK: false,
|
|
Status: status,
|
|
Activities: []*discordgo.Activity{
|
|
{
|
|
Name: "chat",
|
|
Type: discordgo.ActivityTypeListening,
|
|
},
|
|
},
|
|
}
|
|
}
|