diff --git a/chatbot/chatbot.go b/chatbot/chatbot.go index 72e6072..d1bda6d 100644 --- a/chatbot/chatbot.go +++ b/chatbot/chatbot.go @@ -11,6 +11,7 @@ import ( "git.wh64.net/muffin/goMuffin/utils" "github.com/bwmarrin/discordgo" "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" "google.golang.org/genai" ) @@ -84,8 +85,8 @@ func (c *Chatbot) ReloadPrompt() error { } func getMuffinResponse(s *discordgo.Session, question string) (string, error) { - var data []databases.Text var learnData []databases.Learn + var data []databases.Text var result string x := rand.Intn(10) @@ -101,8 +102,15 @@ func getMuffinResponse(s *discordgo.Session, question string) (string, error) { defer muffinCur.Close(context.TODO()) defer learnCur.Close(context.TODO()) - muffinCur.All(context.TODO(), &data) - learnCur.All(context.TODO(), &learnData) + err = muffinCur.All(context.TODO(), &data) + if err != nil { + return "살려주ㅅ세요", err + } + + err = learnCur.All(context.TODO(), &learnData) + if err != nil { + return "살려주ㅅ세요", err + } if x > 2 && len(learnData) != 0 { data := learnData[rand.Intn(len(learnData))] @@ -118,6 +126,7 @@ func getMuffinResponse(s *discordgo.Session, question string) (string, error) { func getAIResponse(s *discordgo.Session, c *Chatbot, user *discordgo.User, question string) (string, error) { var data []databases.Learn + var dbUser databases.User x := rand.Intn(10) @@ -135,7 +144,23 @@ func getAIResponse(s *discordgo.Session, c *Chatbot, user *discordgo.User, quest return fmt.Sprintf("%s\n%s", data.Result, utils.InlineCode(fmt.Sprintf("%s님이 알려주셨어요.", user.Username))), nil } - contents, err := GetMemory(user.ID) + err = databases.Database.Users.FindOne(context.TODO(), databases.User{UserId: user.ID}).Decode(&dbUser) + if err != nil { + return "살려주ㅅ세요", err + } + + err = databases.Database.Chats.FindOne(context.TODO(), databases.Chat{UserId: user.ID}).Err() + if err != nil { + if err == mongo.ErrNoDocuments { + _, err = databases.CreateChat(user.ID, "새로운 채팅") + if err != nil { + return "살려주ㅅ세요", err + } + } + return "살려주ㅅ세요", err + } + + contents, err := GetMemory(dbUser.ChatId) if err != nil { ChatBot.Mode = ChatbotMuffin return "AI에 문제가 생겼ㅇ어요.", err @@ -155,9 +180,10 @@ func getAIResponse(s *discordgo.Session, c *Chatbot, user *discordgo.User, quest UserId: user.ID, Content: question, Answer: resultText, + ChatId: dbUser.ChatId, }) if err != nil { - return "AI에 문제가 생겼ㅇ어요.", err + return "살려주ㅅ세요", err } log.Printf("%s TOKEN: %d", user.ID, result.UsageMetadata.PromptTokenCount) diff --git a/chatbot/memory.go b/chatbot/memory.go index ecc5cdd..6748460 100644 --- a/chatbot/memory.go +++ b/chatbot/memory.go @@ -4,6 +4,7 @@ import ( "context" "git.wh64.net/muffin/goMuffin/databases" + "go.mongodb.org/mongo-driver/v2/bson" "google.golang.org/genai" ) @@ -12,17 +13,20 @@ func SaveMemory(data *databases.Memory) error { return err } -func GetMemory(userId string) ([]*genai.Content, error) { +func GetMemory(chatId bson.ObjectID) ([]*genai.Content, error) { var data []databases.Memory memory := []*genai.Content{} - cur, err := databases.Database.Memory.Find(context.TODO(), databases.User{UserId: userId}) + cur, err := databases.Database.Memory.Find(context.TODO(), databases.User{ChatId: chatId}) if err != nil { return memory, err } - cur.All(context.TODO(), &data) + err = cur.All(context.TODO(), &data) + if err != nil { + return memory, err + } for _, data := range data { memory = append(memory, diff --git a/commands/chat.go b/commands/chat.go index d476356..8c1e66b 100644 --- a/commands/chat.go +++ b/commands/chat.go @@ -1,39 +1,131 @@ package commands import ( + "context" + "fmt" "log" + "strings" "git.wh64.net/muffin/goMuffin/chatbot" + "git.wh64.net/muffin/goMuffin/configs" + "git.wh64.net/muffin/goMuffin/databases" "git.wh64.net/muffin/goMuffin/utils" "github.com/bwmarrin/discordgo" ) +type chatCommandType string + +var ( + chatCommandChatting chatCommandType = "하기" + chatCommandList chatCommandType = "목록" + chatCommandCreate chatCommandType = "생성" +) + var ChatCommand *Command = &Command{ ApplicationCommand: &discordgo.ApplicationCommand{ Name: "대화", - Description: "이 봇이랑 대화해요. (슬래시 커맨드 전용)", + Description: "이 봇이랑 대화해요.", Options: []*discordgo.ApplicationCommandOption{ { - Type: discordgo.ApplicationCommandOptionString, - Name: "내용", - Description: "대화할 내용", - Required: true, + Type: discordgo.ApplicationCommandOptionSubCommand, + Name: string(chatCommandList), + Description: "채팅 목록을 나열해요.", + }, + { + Type: discordgo.ApplicationCommandOptionSubCommand, + Name: string(chatCommandCreate), + Description: "새로운 채팅을 생성해요.", + Options: []*discordgo.ApplicationCommandOption{ + { + Type: discordgo.ApplicationCommandOptionString, + Name: "이름", + Description: "채팅의 이름을 정해요. (25자 이내)", + MaxLength: 25, + Required: true, + }, + }, + }, + { + Type: discordgo.ApplicationCommandOptionSubCommand, + Name: string(chatCommandChatting), + Description: "이 봇이랑 대화해요.", + Options: []*discordgo.ApplicationCommandOption{ + { + Type: discordgo.ApplicationCommandOptionString, + Name: "내용", + Description: "대화할 내용", + Required: true, + }, + }, }, }, }, + Aliases: []string{"채팅"}, DetailedDescription: &DetailedDescription{ - Usage: "/대화 (내용:대화할 내용)", - Examples: []string{"/대화 내용:안녕", "/대화 내용:냠냠"}, + Usage: fmt.Sprintf("%s대화 (목록/생성) [채팅 이름]", configs.Config.Bot.Prefix), + Examples: []string{ + fmt.Sprintf("%s대화 목록", configs.Config.Bot.Prefix), + fmt.Sprintf("%s대화 생성 머핀 냠냠", configs.Config.Bot.Prefix), + }, }, Category: Chatting, RegisterApplicationCommand: true, - RegisterMessageCommand: false, + RegisterMessageCommand: true, Flags: CommandFlagsIsRegistered | CommandFlagsIsBlocked, ChatInputRun: func(ctx *ChatInputContext) error { - i := ctx.Inter - i.DeferReply(&discordgo.InteractionResponseData{}) + ctx.Inter.DeferReply(nil) - str, err := chatbot.ChatBot.GetResponse(i.Member.User, i.Options["내용"].StringValue()) + var cType chatCommandType + var str string + + if opt, ok := ctx.Inter.Options[string(chatCommandChatting)]; ok { + cType = chatCommandChatting + str = opt.Options[0].StringValue() + } else if opt, ok := ctx.Inter.Options[string(chatCommandCreate)]; ok { + cType = chatCommandCreate + str = opt.Options[0].StringValue() + } else if _, ok := ctx.Inter.Options[string(chatCommandList)]; ok { + cType = chatCommandList + } + return chatCommandRun(cType, ctx.Inter, ctx.Inter.User, str) + }, + MessageRun: func(ctx *MsgContext) error { + if len((*ctx.Args)) < 1 { + goto RequiredValue + } + + switch (*ctx.Args)[0] { + case string(chatCommandCreate): + if len((*ctx.Args)) < 2 { + return utils.NewMessageSender(ctx.Msg). + AddComponents(utils.GetErrorContainer(discordgo.TextDisplay{Content: "채팅방의 이름을 정해야해요."})). + SetComponentsV2(true). + SetReply(true). + Send() + } + return chatCommandRun(chatCommandCreate, ctx.Msg, ctx.Msg.Author, string([]rune(strings.Join((*ctx.Args)[1:], " "))[:25])) + case string(chatCommandList): + return chatCommandRun(chatCommandList, ctx.Msg, ctx.Msg.Author, "") + default: + goto RequiredValue + } + + RequiredValue: + return utils.NewMessageSender(ctx.Msg). + AddComponents(utils.GetErrorContainer(discordgo.TextDisplay{Content: "명령어의 첫번째 인자는 `생성`, `목록`중에 하나여야 해요."})). + SetComponentsV2(true). + SetReply(true). + Send() + }, +} + +func chatCommandRun(cType chatCommandType, m any, user *discordgo.User, contentOrName string) error { + switch cType { + // 채팅하기는 슬래시 커맨드만 가능 + case chatCommandChatting: + i := m.(*utils.InteractionCreate) + + str, err := chatbot.ChatBot.GetResponse(user, contentOrName) if err != nil { log.Println(err) i.EditReply(&utils.InteractionEdit{ @@ -42,9 +134,78 @@ var ChatCommand *Command = &Command{ return nil } - result := chatbot.ParseResult(str, ctx.Inter.Session, i) + result := chatbot.ParseResult(str, i.Session, i) return i.EditReply(&utils.InteractionEdit{ Content: &result, }) - }, + case chatCommandCreate: + _, err := databases.CreateChat(user.ID, contentOrName) + if err != nil { + return err + } + return utils.NewMessageSender(m). + AddComponents(utils.GetSuccessContainer(discordgo.TextDisplay{Content: fmt.Sprintf("%s를 생성했어요. 이제 현재 채팅은 %s에요.", contentOrName, contentOrName)})). + SetComponentsV2(true). + SetReply(true). + Send() + case chatCommandList: + var data []databases.Chat + var sections []discordgo.Section + var containers []*discordgo.Container + + cur, err := databases.Database.Chats.Find(context.TODO(), databases.Chat{UserId: user.ID}) + if err != nil { + return err + } + + err = cur.All(context.TODO(), &data) + if err != nil { + return err + } + + if len(data) == 0 { + return utils.NewMessageSender(m). + AddComponents(utils.GetErrorContainer(discordgo.TextDisplay{Content: "채팅이 단 하나도 없어요. 새로운 채팅을 만들거나, 대화를 시작해 채팅을 만들어주세요."})). + SetComponentsV2(true). + SetReply(true). + SetEphemeral(true). + Send() + } + + for i, data := range data { + sections = append(sections, discordgo.Section{ + Accessory: discordgo.Button{ + Label: "선택", + Style: discordgo.SuccessButton, + CustomID: utils.MakeSelectChat(data.Id.Hex(), i+1, user.ID), + }, + Components: []discordgo.MessageComponent{ + discordgo.TextDisplay{ + Content: fmt.Sprintf("%d. %s\n", i+1, data.Name), + }, + }, + }) + } + + textDisplay := discordgo.TextDisplay{Content: fmt.Sprintf("### %s님의 채팅목록", user.GlobalName)} + container := &discordgo.Container{Components: []discordgo.MessageComponent{textDisplay}} + for i, section := range sections { + container.Components = append(container.Components, section, discordgo.Separator{}) + + if (i+1)%10 == 0 { + containers = append(containers, container) + container = &discordgo.Container{Components: []discordgo.MessageComponent{textDisplay}} + continue + } + } + + if len(container.Components) > 1 { + containers = append(containers, container) + } + + return utils.PaginationEmbedBuilder(m). + AddContainers(containers...). + Start() + } + return nil } diff --git a/commands/deleteLearnedData.go b/commands/deleteLearnedData.go index 1dddc2b..ae4fbf2 100644 --- a/commands/deleteLearnedData.go +++ b/commands/deleteLearnedData.go @@ -110,7 +110,6 @@ func deleteLearnedDataRun(m any, command, userId string) error { textDisplay := discordgo.TextDisplay{Content: fmt.Sprintf("### %s 삭제", command)} container := &discordgo.Container{Components: []discordgo.MessageComponent{textDisplay}} - for i, section := range sections { container.Components = append(container.Components, section, discordgo.Separator{}) diff --git a/databases/Chat.go b/databases/Chat.go new file mode 100644 index 0000000..32d9c0e --- /dev/null +++ b/databases/Chat.go @@ -0,0 +1,32 @@ +package databases + +import ( + "context" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" +) + +type Chat struct { + Id bson.ObjectID `bson:"_id,omitempty"` + Name string `bson:"name,omitempty"` + UserId string `bson:"user_id,omitempty"` + CreatedAt time.Time `bson:"created_at,omitempty"` +} + +func CreateChat(userId, name string) (*mongo.InsertOneResult, error) { + createdChat, err := Database.Chats.InsertOne(context.TODO(), Chat{UserId: userId, Name: name, CreatedAt: time.Now()}) + if err != nil { + return nil, err + } + + _, err = Database.Users.UpdateOne(context.TODO(), User{UserId: userId}, bson.D{{ + Key: "$set", + Value: User{ChatId: createdChat.InsertedID.(bson.ObjectID)}, + }}) + if err != nil { + return nil, err + } + return createdChat, nil +} diff --git a/databases/User.go b/databases/User.go index 855cb9a..d99a86e 100644 --- a/databases/User.go +++ b/databases/User.go @@ -12,6 +12,7 @@ type User struct { UserId string `bson:"user_id,omitempty"` Blocked bool `bson:"blocked,omitempty"` BlockedReason string `bson:"blocked_reason,omitempty"` + ChatId bson.ObjectID `bson:"chat_id,omitempty"` CreatedAt time.Time `bson:"created_at,omitempty"` } diff --git a/databases/database.go b/databases/database.go index fe7da03..be6421f 100644 --- a/databases/database.go +++ b/databases/database.go @@ -14,6 +14,7 @@ type MuffinDatabase struct { Texts *mongo.Collection Memory *mongo.Collection Users *mongo.Collection + Chats *mongo.Collection } var Database *MuffinDatabase @@ -38,5 +39,6 @@ func Connect() (*MuffinDatabase, error) { Texts: client.Database(configs.Config.Database.Name).Collection("text"), Memory: client.Database(configs.Config.Database.Name).Collection("memory"), Users: client.Database(configs.Config.Database.Name).Collection("user"), + Chats: client.Database(configs.Config.Database.Name).Collection("chat"), }, nil } diff --git a/utils/customIds.go b/utils/customIds.go index 3e1f894..86a5e4a 100644 --- a/utils/customIds.go +++ b/utils/customIds.go @@ -22,6 +22,8 @@ const ( DeregisterAgree = "#muffin/deregister/agree@" DeregisterDisagree = "#muffin/deregister/disagree@" + + SelectChat = "#muffin/chat/select@" ) func MakeDeleteLearnedData(id string, number int, userId string) string { @@ -117,3 +119,7 @@ func GetDeregisterUserId(customId string) string { return customId } } + +func MakeSelectChat(id string, number int, userId string) string { + return fmt.Sprintf("%sid=%s&no=%d&user_id=%s", DeleteLearnedData, id, number, userId) +}