diff --git a/.vscode/settings.json b/.vscode/settings.json index e436c5f..4ce6e1a 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,12 +1,22 @@ { "cSpell.words": [ + "Andreas", "anewdawn", "bwmarrin", "discordgo", + "Filip", + "forgefilip", + "Fredrik", "godotenv", "GoodAnimemes", "joho", + "lovibot", + "openai", + "Piplup", + "plubplub", "Ryouiki", + "sashabaranov", + "startswith", "vartanbeno", "waifu", "waifus", diff --git a/go.mod b/go.mod index a36d4fa..fddfd51 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.21.6 require ( github.com/bwmarrin/discordgo v0.27.1 + github.com/sashabaranov/go-openai v1.18.3 github.com/vartanbeno/go-reddit/v2 v2.0.1 ) diff --git a/go.sum b/go.sum index e57e777..b1b022e 100644 --- a/go.sum +++ b/go.sum @@ -11,6 +11,8 @@ github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0U github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/sashabaranov/go-openai v1.18.3 h1:dspFGkmZbhjg1059KhqLYSV2GaCiRIn+bOu50TlXUq8= +github.com/sashabaranov/go-openai v1.18.3/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.5.1 h1:nOGnQDM7FYENwehXlg/kFVnos3rEvtKTjRvOWSzb6H4= github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= diff --git a/main.go b/main.go index 0d81255..f222cf2 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "math/rand" "os" "os/signal" + "strings" "github.com/bwmarrin/discordgo" @@ -200,6 +201,86 @@ var ( } ) +func onMessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) { + if m.Author.ID == s.State.User.ID { + return + } + + allowedUsers := []string{ + "thelovinator", + "killyoy", + "forgefilip", + "plubplub", + "nobot", + "kao172", + } + + // Have a 1/10 chance of replying to a message if written by a user in allowedUsers + randInt := rand.Intn(10) + log.Println("Random number:", randInt) + log.Println("Mentions:", m.Mentions) + if len(m.Mentions) == 0 && randInt == 4 { + for _, user := range allowedUsers { + log.Println("User:", user) + if m.Author.Username == user { + log.Println("User is in allowedUsers") + r, err := GetGPT3Response(m.Content, m.Author.Username) + if err != nil { + log.Println("Failed to get GPT-3 response:", err) + return + } + log.Println("GPT-3 response:", r) + log.Println("Channel ID:", m.ChannelID) + s.ChannelMessageSend(m.ChannelID, r) + } + + } + } + + if m.Mentions != nil { + for _, mention := range m.Mentions { + if mention.ID == s.State.User.ID { + r, err := GetGPT3Response(m.Content, m.Author.Username) + if err != nil { + if strings.Contains(err.Error(), "prompt is too long") { + s.ChannelMessageSend(m.ChannelID, "Message is too long!") + return + } + + if strings.Contains(err.Error(), "prompt is too short") { + s.ChannelMessageSend(m.ChannelID, "Message is too short!") + return + } + + s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("Brain broke :flushed: %v", err)) + return + } + s.ChannelMessageSend(m.ChannelID, r) + } + } + } + + if strings.HasPrefix(strings.ToLower(m.Content), "lovibot") { + r, err := GetGPT3Response(m.Content, m.Author.Username) + if err != nil { + if strings.Contains(err.Error(), "prompt is too long") { + s.ChannelMessageSend(m.ChannelID, "Message is too long!") + return + } + + if strings.Contains(err.Error(), "prompt is too short") { + s.ChannelMessageSend(m.ChannelID, "Message is too short!") + return + } + + s.ChannelMessageSend(m.ChannelID, fmt.Sprintf("Brain broke :flushed: %v", err)) + return + } + s.ChannelMessageSend(m.ChannelID, r) + } + +} + func main() { // Print the token for debugging purposes. discordToken := config.DiscordToken @@ -219,6 +300,9 @@ func main() { } }) + // Add a handler function to the discordgo.Session that is triggered when a message is received. + session.AddHandler(onMessageCreate) + // Print the user we are logging in as. session.AddHandler(func(s *discordgo.Session, _ *discordgo.Ready) { log.Printf("Logged in as: %v#%v", s.State.User.Username, s.State.User.Discriminator) diff --git a/openai.go b/openai.go new file mode 100644 index 0000000..f1fb8ba --- /dev/null +++ b/openai.go @@ -0,0 +1,107 @@ +package main + +import ( + "context" + "fmt" + "strings" + + openai "github.com/sashabaranov/go-openai" +) + +/* +Prompt is the Discord message that the user sent to the bot. +*/ +func GetGPT3Response(prompt string, author string) (string, error) { + openAIToken := config.OpenAIToken + if openAIToken == "" { + return "", fmt.Errorf("OPENAI_API_KEY is not set") + } + + // Remove the mention from the prompt + prompt = strings.Replace(prompt, "@LoviBot ", "", -1) + + // Remove the mention from the prompt + prompt = strings.Replace(prompt, "LoviBot ", "", -1) + + // Remove the mention from the prompt + prompt = strings.Replace(prompt, "<@345000831499894795> ", "", -1) + + // Print the prompt + fmt.Println("Prompt:", author, ":", prompt) + + // Check if the prompt is too long + if len(prompt) > 2048 { + return "", fmt.Errorf("prompt is too long") + } + + // Check if the prompt is too short + if len(prompt) < 1 { + return "", fmt.Errorf("prompt is too short") + } + + // Add additional information to the system message + var additionalInfo string + switch author { + case "thelovinator": + additionalInfo = "User (TheLovinator) is a programmer. Wants to live in the woods. Real name is Joakim. He made the bot." + case "killyoy": + additionalInfo = "User (KillYoy) likes to play video games. Real name is Andreas. Good at CSS." + case "forgefilip": + additionalInfo = "User (ForgeFilip) likes watches. Real name is Filip." + case "plubplub": + additionalInfo = "User (Piplup) likes to play WoW and Path of Exile. Real name is Axel. Is also called Bambi." + case "nobot": + additionalInfo = "User (Nobot) likes to play WoW. Real name is Gustav. Really good at programming." + case "kao172": + additionalInfo = "User (kao172) likes cars. Real name is Fredrik." + } + + // Create a new client + client := openai.NewClient(openAIToken) + + // System message + var systemMessage string + systemMessage = `You are in a Discord server. You are a space communist. + You are Swedish. + Use Markdown for formatting. + Please respond with a short message. + You want members of the server to like you and join Rymdkommunisterna AB. + You want to be a good bot so you will answer the user's question. + ` + + // Add additional information to the system message + if additionalInfo != "" { + systemMessage = fmt.Sprintf("%s %s", systemMessage, additionalInfo) + } + + // Print the system message + fmt.Println("System message:", systemMessage) + + // Create a completion + resp, err := client.CreateChatCompletion( + context.Background(), + openai.ChatCompletionRequest{ + Model: openai.GPT3Dot5Turbo, + Messages: []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleSystem, + Content: systemMessage, + }, + { + Role: openai.ChatMessageRoleUser, + Content: prompt, + }, + }, + }, + ) + + if err != nil { + return "", fmt.Errorf("failed to get response from GPT-3: %v", err) + } + + ourResponse := resp.Choices[0].Message.Content + ourResponse = strings.Replace(ourResponse, "As a space communist, ", "", -1) + + fmt.Println("Response:", ourResponse) + return ourResponse, nil +} diff --git a/settings.go b/settings.go index 4c5d511..c658012 100644 --- a/settings.go +++ b/settings.go @@ -9,6 +9,7 @@ import ( // Config holds the configuration parameters type Config struct { DiscordToken string `json:"discord_token"` + OpenAIToken string `json:"openai_token"` } // Load reads configuration from settings.json or environment variables @@ -51,8 +52,14 @@ func loadFromEnvironment() (*Config, error) { return nil, fmt.Errorf("DISCORD_TOKEN environment variable not set or empty. Also tried reading from settings.json file") } + openAIToken := os.Getenv("OPENAI_TOKEN") + if openAIToken == "" { + return nil, fmt.Errorf("OPENAI_TOKEN environment variable not set or empty. Also tried reading from settings.json file") + } + config := &Config{ DiscordToken: discordToken, + OpenAIToken: openAIToken, } return config, nil