feat: add chat command

This commit is contained in:
Siwoo Jeon 2025-07-11 18:33:26 +09:00
parent 76cacdd660
commit 6a78106136
Signed by: migan
GPG key ID: 036E9A8C5E8E48DA
8 changed files with 253 additions and 22 deletions

View file

@ -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)

View file

@ -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,

View file

@ -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
}

View file

@ -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{})

32
databases/Chat.go Normal file
View file

@ -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
}

View file

@ -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"`
}

View file

@ -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
}

View file

@ -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)
}