goMuffin/chatbot/chatbot.go

package chatbot

import (
	"context"
	"fmt"
	"math/rand"
	"sync"

	"git.wh64.net/muffin/goMuffin/configs"
	"git.wh64.net/muffin/goMuffin/databases"
	"git.wh64.net/muffin/goMuffin/utils"
	"github.com/bwmarrin/discordgo"
	"go.mongodb.org/mongo-driver/v2/bson"
	"google.golang.org/genai"
)

type Chatbot struct {
	Mode         ChatbotMode
	Gemini       *genai.Client
	systemPrompt string
	s            *discordgo.Session
}

var ChatBot *Chatbot

func New(s *discordgo.Session) error {
	gemini, err := genai.NewClient(context.TODO(), &genai.ClientConfig{
		APIKey:  configs.Config.Chatbot.Gemini.Token,
		Backend: genai.BackendGeminiAPI,
	})
	if err != nil {
		return err
	}

	ChatBot = &Chatbot{
		Mode:   ChatbotAI,
		Gemini: gemini,
		s:      s,
	}

	prompt, err := loadPrompt()
	if err != nil {
		return err
	}

	ChatBot.systemPrompt = prompt
	return nil
}

func (c *Chatbot) SetMode(mode ChatbotMode) *Chatbot {
	c.Mode = mode
	return c
}

func (c *Chatbot) SwitchMode() *Chatbot {
	switch c.Mode {
	case ChatbotAI:
		c.SetMode(ChatbotMuffin)
	case ChatbotMuffin:
		c.SetMode(ChatbotAI)
	}
	return c
}

func (c *Chatbot) ModeString() string {
	switch c.Mode {
	case ChatbotAI:
		return "AI모드"
	case ChatbotMuffin:
		return "머핀 모드"
	default:
		return "알 수 없음"
	}
}

func (c *Chatbot) ReloadPrompt() error {
	prompt, err := loadPrompt()
	if err != nil {
		return err
	}

	c.systemPrompt = prompt
	return nil
}

func getMuffinResponse(s *discordgo.Session, question string) (string, error) {
	var data []databases.Text
	var learnData []databases.Learn
	var result string
	var wg sync.WaitGroup

	ch1 := make(chan error)
	ch2 := make(chan error)
	x := rand.Intn(10)

	wg.Add(2)

	// 머핀 데이터
	go func() {
		cur, err := databases.Database.Texts.Find(context.TODO(), bson.D{{Key: "persona", Value: "muffin"}})
		if err != nil {
			ch1 <- err
		}

		defer cur.Close(context.TODO())

		cur.All(context.TODO(), &data)
		ch1 <- nil
		wg.Done()
	}()

	// 지식 데이터
	go func() {
		cur, err := databases.Database.Learns.Find(context.TODO(), bson.D{{Key: "command", Value: question}})
		if err != nil {
			ch2 <- err
		}

		defer cur.Close(context.TODO())

		cur.All(context.TODO(), &learnData)
		ch2 <- nil
		wg.Done()
	}()

	wg.Wait()
	select {
	case err := <-ch1:
		if err != nil {
			return "에러 발생", fmt.Errorf("muffin data error\n%s", err.Error())
		}
	case err := <-ch2:
		if err != nil {
			return "에러 발생", fmt.Errorf("learn data error\n%s", err.Error())
		}
	}

	if x > 2 && len(learnData) != 0 {
		data := learnData[rand.Intn(len(learnData))]
		user, _ := s.User(data.UserId)

		result =
			fmt.Sprintf("%s\n%s", data.Result, utils.InlineCode(fmt.Sprintf("%s님이 알려주셨어요.", user.Username)))
	} else {
		result = data[rand.Intn(len(data))].Text
	}
	return result, nil
}

func getAIResponse(c *Chatbot, user *discordgo.User, question string) (string, error) {
	contents, err := GetMemory(user.ID)
	if err != nil {
		ChatBot.Mode = ChatbotMuffin
		return "AI에 문제가 생겼ㅇ어요.", err
	}

	contents = append(contents, genai.NewContentFromText(question, genai.RoleUser))
	result, err := ChatBot.Gemini.Models.GenerateContent(context.TODO(), configs.Config.Chatbot.Gemini.Model, contents, &genai.GenerateContentConfig{
		SystemInstruction: genai.NewContentFromText(makePrompt(c.systemPrompt, user), genai.RoleUser),
	})
	if err != nil {
		ChatBot.Mode = ChatbotMuffin
		return "AI에 문제가 생겼ㅇ어요.", err
	}

	resultText := result.Text()
	err = SaveMemory(&databases.InsertMemory{
		UserId:  user.ID,
		Content: question,
		Answer:  resultText,
	})
	if err != nil {
		return "", err
	}

	return resultText, nil
}

func (c *Chatbot) GetResponse(user *discordgo.User, question string) (string, error) {
	switch c.Mode {
	case ChatbotMuffin:
		return getMuffinResponse(c.s, question)
	default:
		return getAIResponse(c, user, question)
	}
}