ai work

plugins/rp.go

@@ -0,0 +1,229 @@
+package plugins
+
+import (
+	"database/sql"
+	"errors"
+	"fmt"
+	"kurumibot/database/mdb"
+	"kurumibot/database/psql"
+	"kurumibot/database/red"
+	"kurumibot/laniakea"
+	"kurumibot/utils/ai"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/google/uuid"
+)
+
+func RegisterRP(bot *laniakea.Bot) {
+	rp := laniakea.NewPlugin("RP")
+	rp = rp.Command(selectWaifu, "rpwaifu", "рпвайфу")
+	rp = rp.Payload(selectWaifu, "rp.selwaifu")
+	rp = rp.Command(rpPresetsList, "rpplist")
+	rp = rp.Command(rpPresetSet, "rppset")
+	rp = rp.Command(newChat, "newchat")
+	rp = rp.Command(generate, "g", "gen", "г")
+
+	bot.AddPlugins(rp.Build())
+}
+
+func selectWaifu(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
+	waifuId, err := strconv.Atoi(ctx.Args[0])
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	err = red.RPSetSelectedWaifu(db, ctx.FromID, waifuId)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+
+	ctx.Answer(fmt.Sprintf("Была выбрана вайфу %d", waifuId))
+}
+func rpPresetsList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
+	rep := psql.NewRPRepository(db.PostgresSQL)
+	presets, err := rep.GetAllPresets()
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	out := make([]string, len(presets))
+	for i, preset := range presets {
+		out[i] = fmt.Sprintf("%s) *%s*\n%s", preset.ID, preset.Name, preset.Description)
+	}
+	ctx.Answer(strings.Join(out, "\n"))
+}
+func rpPresetSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
+	if len(ctx.Args) == 0 || ctx.Args[0] == "" {
+		return
+	}
+	presetId := ctx.Args[0]
+	rep := psql.NewRPRepository(db.PostgresSQL)
+	user, err := rep.GetOrCreateUser(int64(ctx.FromID))
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	preset, err := rep.UpdateUserPreset(user, presetId)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			ctx.Answer("Данный пресет не найден")
+		} else {
+			ctx.Error(err)
+		}
+		return
+	}
+	ctx.Answer(fmt.Sprintf("Был выбран пресет %s", preset.Name))
+}
+func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
+	waifuId := red.RPGetSelectedWaifu(db, ctx.FromID)
+	if waifuId == 0 {
+		ctx.Answer("Не выбрана вайфу")
+		return
+	}
+	chatId := uuid.New()
+	err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId.String())
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+
+	prompt := strings.Join(ctx.Args[1:], " ")
+	err = red.RPSetChatPrompt(db, ctx.FromID, waifuId, prompt)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	err = red.RPSetCounter(db, ctx.FromID, waifuId, 0)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	ctx.Answer("Был создан новый чат.")
+}
+
+func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
+	waifuId := red.RPGetSelectedWaifu(db, ctx.FromID)
+	if waifuId == 0 {
+		ctx.Answer("Не выбрана вайфу")
+		return
+	}
+	chatId := red.RPGetChatId(db, ctx.FromID, waifuId)
+	if chatId == "" {
+		chatId = uuid.New().String()
+		err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId)
+		if err != nil {
+			ctx.Error(err)
+			return
+		}
+	}
+	waifu, err := psql.GetWaifuById(waifuId)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	rpRep := psql.NewRPRepository(db.PostgresSQL)
+	rpUser, err := rpRep.GetUser(int64(ctx.FromID))
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	preset, err := rpRep.GetUserPreset(rpUser)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+
+	systemPrompt := ai.Message{
+		Role: "system",
+		Content: fmt.Sprintf(
+			"%s %s %s",
+			ai.FormatPrompt(preset.PreHistory, waifu.Name, ctx.Msg.From.FirstName),
+			fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt),
+			red.RPGetChatPrompt(db, ctx.FromID, waifuId),
+		),
+	}
+
+	history, err := mdb.GetChatHistory(db, chatId)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	messages := make([]ai.Message, 0)
+	for _, m := range history {
+		messages = append(messages, ai.Message{
+			Role:    m.Role,
+			Content: m.Message,
+		})
+	}
+
+	api := ai.NewOpenAIAPI(ai.CosmoRPUrl, os.Getenv("PAWAN_KEY"), "cosmorp-2.5")
+	userMessage := strings.Join(ctx.Args, " ")
+	messages = append(messages, ai.Message{
+		Role: "system", Content: ai.FormatPrompt(preset.PostHistory, waifu.Name, ctx.Msg.From.FirstName),
+	}, ai.Message{
+		Role: "user", Content: userMessage,
+	})
+	err = mdb.UpdateChatHistory(db, chatId, "user", userMessage)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+
+	m := ctx.Answer("Генерация запущена...")
+	res, err := api.CreateCompletion(ai.CreateCompletionReq{
+		Messages:    append([]ai.Message{systemPrompt}, messages...),
+		Verbosity:   "low",
+		Temperature: 0.7,
+	})
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	response := make([]string, 0)
+	for _, choice := range res.Choices {
+		m := choice.Message
+		response = append(response, m.Content)
+		err = mdb.UpdateChatHistory(db, chatId, m.Role, m.Content)
+	}
+	rpUser.UsedTokens = rpUser.UsedTokens + res.Usage.TotalTokens
+	err = rpRep.UpdateUser(rpUser)
+	if err != nil {
+		ctx.Error(err)
+		return
+	}
+	m.Delete()
+	ctx.Answer(laniakea.EscapeMarkdown(strings.Join(response, "\n")))
+
+	counter := red.RPGetCounter(db, ctx.FromID, waifuId) + 1
+	if counter == 5 {
+		res, err = api.CreateCompletion(ai.CreateCompletionReq{
+			Messages: append(messages, ai.Message{
+				Role:    "user",
+				Content: "Сделай краткую выжимку нашего диалога",
+			}),
+		})
+		if err != nil {
+			ctx.Error(err)
+		}
+		chatId = uuid.New().String()
+		err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId)
+		if err != nil {
+			ctx.Error(err)
+			return
+		}
+		err = mdb.UpdateChatHistory(db, chatId, "assistant", res.Choices[0].Message.Content)
+		err = red.RPSetCounter(db, ctx.FromID, waifuId, 0)
+		if err != nil {
+			ctx.Error(err)
+			return
+		}
+	} else {
+		err = red.RPSetCounter(db, ctx.FromID, waifuId, counter)
+		if err != nil {
+			ctx.Error(err)
+		}
+	}
+}